mcp-go icon indicating copy to clipboard operation
mcp-go copied to clipboard

[proposal] Optimize hooks

Open dugenkui03 opened this issue 1 year ago • 1 comments

background

The definition of hooks are inconvenient, because the methods of hooks is scattered. If I want to define two kind of hooks for two different scenarios, such as logging and monitoring, I can't manage two kind of hook struct, and can only create some functions.

Moreover, I believe that hooks are not just for obtaining and recording execution information; they can also influence the execution process, like parameter validation. However, the current hook definition does not support such operations.

change

I suggest that:

  1. Add Hook interface and ChainHooks struct( ChainHooks also implements the Hook interface)
  2. Add error result on some method, if return error, interrupt the process
  3. Add HookContext to save hook context

type Hook interface {
    CreateHookContext() HookContext
    GetHookContext() HookContext
    ......
    OnBeforeListTools(ctx context.Context, id any, message *mcp.ListToolsRequest)
    OnAfterListTools(ctx context.Context, id any, message *mcp.ListToolsRequest, result *mcp.ListToolsResult)
    // developer can use this method validate tool call input
    OnBeforeCallTool(ctx context.Context, id any, message *mcp.CallToolRequest) error
    OnAfterCallTool(ctx context.Context, id any, message *mcp.CallToolRequest, result *mcp.CallToolResult)
}

// ChainHooks also implements the Hook interface.
type ChainHooks struct {
    Hooks []Hook
}

func (c *ChainHooks) OnBeforeListTools(ctx context.Context, id any, message *mcp.ListToolsRequest) {
    for _, hook := range c.Hooks {
        hook.OnBeforeListTools(ctx, id, message)
    }
}

func (c *ChainHooks) OnAfterListTools(ctx context.Context, id any, message *mcp.ListToolsRequest, result *mcp.ListToolsResult) {
    for _, hook := range c.Hooks {
        hook.OnAfterListTools(ctx, id, message, result)
    }
}

func (c *ChainHooks) OnBeforeCallTool(ctx context.Context, id any, message *mcp.CallToolRequest) error {
    for _, hook := range c.Hooks {
        err:= hook.OnBeforeCallTool(ctx, id, message)
        if err != nil {
            return err
        }
    }
}

func (c *ChainHooks) OnAfterCallTool(ctx context.Context, id any, message *mcp.CallToolRequest, result *mcp.CallToolResult) {
    for _, hook := range c.Hooks {
        hook.OnAfterCallTool(ctx, id, message, result)
    }
}
type HookContext interface{
	IsHookContext()
}

type ChainedHookContext struct{
	HookContexts []HookContext
}

dugenkui03 avatar Apr 16 '25 13:04 dugenkui03

About HookContext

HandleMessage will invoke hooks#CreateHookContext to create HookContext to save request context.

Also, you can always return same object to do some statistical thing:


type globalHookContext struct {
    // ....... some statistical fields
}

func (*globalHookContext) IsHookContext() {}



var globalHookCtx = &globalHookContext{}

func (*globalHook) CreateHookContext() HookContext {
	// All request share one hook context
	return globalHookCtx
}

dugenkui03 avatar Apr 16 '25 19:04 dugenkui03