diff --git a/agent/agent_test.go b/agent/agent_test.go index c1bef9e8..21f2434f 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -6,6 +6,7 @@ import ( "context" "errors" "iter" + "reflect" "slices" "sync" "testing" @@ -21,6 +22,106 @@ type stubTool struct { name string } +func TestFunctionInvocationMiddleware_Composition(t *testing.T) { + toolFailure := errors.New("tool failed") + for _, tc := range []struct { + name string + short bool + toolError error + }{ + {name: "arguments and result replacement"}, + {name: "error propagation", toolError: toolFailure}, + {name: "short circuit", short: true}, + } { + t.Run(tc.name, func(t *testing.T) { + var order []string + fn := functool.MustNew(functool.Config{Name: "lookup", Description: "Look up a value"}, func(_ context.Context, args struct { + Value string `json:"value"` + }, + ) (string, error) { + order = append(order, "tool") + if args.Value != "changed" { + t.Errorf("tool argument = %q, want changed", args.Value) + } + return args.Value, tc.toolError + }) + first := agent.FunctionInvocationMiddleware(func(ctx context.Context, invocation *agent.FunctionInvocationContext, next agent.FunctionInvocationFunc) (any, error) { + order = append(order, "first before") + if invocation.Function != fn || invocation.Arguments != `{"value":"original"}` || invocation.CallID != "" { + t.Errorf("unexpected direct invocation: %#v", invocation) + } + if tc.short { + return "cached", nil + } + invocation.Arguments = `{"value":"changed"}` + result, err := next(ctx, invocation) + order = append(order, "first after") + if err != nil { + return nil, err + } + return "wrapped " + result.(string), nil + }) + second := agent.FunctionInvocationMiddleware(func(ctx context.Context, invocation *agent.FunctionInvocationContext, next agent.FunctionInvocationFunc) (any, error) { + order = append(order, "second before") + result, err := next(ctx, invocation) + order = append(order, "second after") + return result, err + }) + other := stubTool{name: "hosted"} + options := []agent.Option{agent.WithTool(fn), agent.WithTool(other), agent.WithTool(nil)} + provider := agent.NewContextProvider(agent.ContextProviderConfig{SourceID: "tools", Provide: func(context.Context, agent.InvokingContext) ([]*message.Message, []agent.Option, error) { + return nil, options, nil + }}) + run := func(ctx context.Context, _ []*message.Message, opts ...agent.Option) iter.Seq2[*agent.ResponseUpdate, error] { + return func(yield func(*agent.ResponseUpdate, error) bool) { + tools := slices.Collect(agent.AllOptions(opts, agent.WithTool)) + if len(tools) != 2 || tools[1] != other { + t.Fatalf("unexpected tools: %v", tools) + } + wrapped := tools[0].(tool.FuncTool) + if wrapped.Name() != fn.Name() || wrapped.Description() != fn.Description() || + !reflect.DeepEqual(wrapped.Schema(), fn.Schema()) || !reflect.DeepEqual(wrapped.ReturnSchema(), fn.ReturnSchema()) { + t.Error("wrapper changed tool metadata") + } + result, err := wrapped.Call(ctx, `{"value":"original"}`) + if err != nil { + yield(nil, err) + return + } + want := "wrapped changed" + if tc.short { + want = "cached" + } + if result != want { + t.Errorf("result = %v, want %q", result, want) + } + yield(&agent.ResponseUpdate{Role: message.RoleAssistant, Contents: message.Contents{&message.TextContent{Text: "done"}}}, nil) + } + } + var nilMiddleware agent.FunctionInvocationMiddleware + a := agent.New(agent.ProviderConfig{Run: run, Middlewares: []agent.Middleware{first, nilMiddleware, second}}, agent.Config{ + ContextProviders: []agent.ContextProvider{provider}, + }) + for range 2 { + order = nil + if _, err := a.RunText(t.Context(), "start").Collect(); !errors.Is(err, tc.toolError) { + t.Fatalf("error = %v, want %v", err, tc.toolError) + } + wantOrder := []string{"first before", "second before", "tool", "second after", "first after"} + if tc.short { + wantOrder = []string{"first before"} + } + if !slices.Equal(order, wantOrder) { + t.Errorf("callback order = %v, want %v", order, wantOrder) + } + } + if original, _ := agent.GetOption(options[:1], agent.WithTool); original != fn { + t.Error("middleware mutated the context provider's options") + } + }) + } +} + func (t stubTool) Name() string { return t.name } diff --git a/agent/harness/toolautocall/autocall.go b/agent/harness/toolautocall/autocall.go index a0611b0a..786d0a2b 100644 --- a/agent/harness/toolautocall/autocall.go +++ b/agent/harness/toolautocall/autocall.go @@ -20,6 +20,8 @@ import ( "github.com/microsoft/agent-framework-go/agent" "github.com/microsoft/agent-framework-go/internal/otelx" "github.com/microsoft/agent-framework-go/internal/slogx" + "github.com/microsoft/agent-framework-go/internal/toolcontext" + "github.com/microsoft/agent-framework-go/internal/toolmiddleware" "github.com/microsoft/agent-framework-go/message" "github.com/microsoft/agent-framework-go/tool" @@ -64,6 +66,7 @@ type Config struct { // provider requests. Request tools supplied through agent options take // precedence; this collection is consulted afterward, which is useful when the // provider is already configured with tool declarations out of band. + // Function invocation middleware registered before this middleware also wraps these tools. AdditionalTools []tool.Tool // IncludeDetailedErrors controls whether tool error details are included in @@ -215,7 +218,7 @@ func (f *autocall) Run(next agent.RunFunc, ctx context.Context, messages []*mess } messagesCloned = messagesCloned || changed } - tools, _ := f.createToolsMap(agent.AllOptions(opts, agent.WithTool)) + tools, _ := f.createToolsMap(opts) // This is a synthetic ID since we're generating the tool messages instead of getting them from // the underlying provider. When emitting the streamed chunks, it's perfectly valid for us to @@ -279,7 +282,7 @@ func (f *autocall) Run(next agent.RunFunc, ctx context.Context, messages []*mess f.logger.Debug(ctx, "reached maximum iteration count; stopping function invocation loop", "maximumIterationsPerRequest", f.maximumIterationsPerRequest) opts = prepareOptionsForLastIteration(opts) } - tools, requiresApproval := f.createToolsMap(agent.AllOptions(opts, agent.WithTool)) + tools, requiresApproval := f.createToolsMap(opts) // Reset slice without reallocating. updates = updates[:0] @@ -619,7 +622,7 @@ func (f *autocall) shouldTerminateLoopBasedOnHandleableFunctions(ctx context.Con return false } -func (f *autocall) createToolsMap(tools iter.Seq[tool.Tool]) (mtools map[string]tool.SchemaTool, anyRequiredApproval bool) { +func (f *autocall) createToolsMap(opts []agent.Option) (mtools map[string]tool.SchemaTool, anyRequiredApproval bool) { fn := func(t tool.Tool) { if !anyRequiredApproval { if approval, ok := t.(tool.ApprovalRequiredTool); ok && approval.ApprovalRequired() { @@ -638,10 +641,18 @@ func (f *autocall) createToolsMap(tools iter.Seq[tool.Tool]) (mtools map[string] } mtools[declaration.Name()] = declaration } - for t := range tools { + for t := range agent.AllOptions(opts, agent.WithTool) { fn(t) } for _, t := range f.additionalTools { + if function, ok := t.(tool.FuncTool); ok { + for _, opt := range opts { + if wrap, ok := opt.(toolmiddleware.Wrapper); ok { + function = wrap(function) + } + } + t = function + } fn(t) } return mtools, anyRequiredApproval @@ -1154,6 +1165,7 @@ func (f *autocall) processFunctionCall(ctx context.Context, tools map[string]too } f.logger.Debug(ctx, "calling function", "funcName", funcCall.Name, slogx.SensitiveData("arguments", funcCall.Arguments)) start := time.Now() + ctx = toolcontext.WithCallID(ctx, funcCall.CallID) ctx, span := startToolSpan(ctx, funcCall, declaration) if span != nil { defer span.End() diff --git a/agent/harness/toolautocall/autocall_approval_test.go b/agent/harness/toolautocall/autocall_approval_test.go index 8881fee1..10b1874b 100644 --- a/agent/harness/toolautocall/autocall_approval_test.go +++ b/agent/harness/toolautocall/autocall_approval_test.go @@ -26,6 +26,111 @@ func expectedMessages(t *testing.T, expected ...*message.Message) func(context.C } } +func TestFunctionInvoking_InvocationIdentityAfterApproval(t *testing.T) { + for _, tc := range []struct { + name string + approved bool + shortCircuit bool + additional bool + wantTools int + wantWrappers int + }{ + {name: "approved", approved: true, wantTools: 1, wantWrappers: 1}, + {name: "denied"}, + {name: "approved short circuit", approved: true, shortCircuit: true, wantWrappers: 1}, + {name: "additional tool approved", additional: true, approved: true, wantTools: 1, wantWrappers: 1}, + {name: "additional tool denied", additional: true}, + {name: "additional tool approved short circuit", additional: true, approved: true, shortCircuit: true, wantWrappers: 1}, + } { + t.Run(tc.name, func(t *testing.T) { + var observedIDs []string + var toolCalls int + testTool := tool.ApprovalRequiredFunc(functool.MustNew(functool.Config{Name: "lookup"}, func(ctx context.Context, _ struct{}) (string, error) { + toolCalls++ + invocation, ok := tool.InvocationFromContext(ctx) + if !ok || invocation.CallID != "call-1" { + t.Errorf("handler identity = %#v, %v; want call-1", invocation, ok) + } + return "done", nil + })) + observer := agent.FunctionInvocationMiddleware(func(ctx context.Context, invocation *agent.FunctionInvocationContext, next agent.FunctionInvocationFunc) (any, error) { + observedIDs = append(observedIDs, invocation.CallID) + if tc.shortCircuit { + return "cached", nil + } + return next(ctx, invocation) + }) + var providerResumed bool + runner := &agenttest.Runner{Responses: agenttest.NewResponseBuilder(). + AddFunctionCall("call-1", "lookup", `{}`). + NewTurn(func(context.Context, []*message.Message, ...agent.Option) { providerResumed = true }). + AddText("done").Build()} + cfg := toolautocall.Config{} + tools := []tool.Tool{testTool} + if tc.additional { + cfg.AdditionalTools = tools + tools = nil + } + a := agent.New(agent.ProviderConfig{ + Run: runner.Run, Middlewares: []agent.Middleware{toolautocall.New(cfg)}, + }, agent.Config{Tools: tools, Middlewares: []agent.Middleware{observer}}) + session := &agent.Session{} + var request *message.ToolApprovalRequestContent + for update, err := range a.RunText(t.Context(), "start", agent.WithSession(session)) { + if err != nil { + t.Fatal(err) + } + for _, content := range update.Contents { + if approval, ok := content.(*message.ToolApprovalRequestContent); ok { + request = approval + } + } + } + if request == nil || toolCalls != 0 || len(observedIDs) != 0 { + t.Fatalf("expected approval before execution; request=%v tool calls=%d wrapper calls=%d", request, toolCalls, len(observedIDs)) + } + call := request.ToolCall.(*message.FunctionCallContent) + if call.CallID != "call-1" { + t.Fatalf("approval call ID = %q, want call-1", call.CallID) + } + data, err := json.Marshal(session) + if err != nil { + t.Fatal(err) + } + var restored agent.Session + if err := json.Unmarshal(data, &restored); err != nil { + t.Fatal(err) + } + var resultIDs []string + for update, err := range a.RunMessage(t.Context(), message.New(request.CreateResponse(tc.approved, "")), agent.WithSession(&restored)) { + if err != nil { + t.Fatal(err) + } + for _, content := range update.Contents { + if result, ok := content.(*message.FunctionResultContent); ok { + resultIDs = append(resultIDs, result.CallID) + if tc.shortCircuit && result.Result != "cached" { + t.Errorf("short-circuit result = %v, want cached", result.Result) + } + } + } + } + if toolCalls != tc.wantTools || len(observedIDs) != tc.wantWrappers { + t.Fatalf("tool calls=%d wrapper calls=%d, want %d and %d", toolCalls, len(observedIDs), tc.wantTools, tc.wantWrappers) + } + if tc.approved && observedIDs[0] != call.CallID { + t.Errorf("resumed call ID = %q, want %q", observedIDs[0], call.CallID) + } + if !providerResumed { + t.Error("provider did not receive the result after approval handling") + } + if len(resultIDs) != 1 || resultIDs[0] != call.CallID { + t.Errorf("result IDs = %v, want [%s]", resultIDs, call.CallID) + } + }) + } +} + // invokeAndAssertApproval is the helper for approval tests func invokeAndAssertApproval(t *testing.T, tools []tool.Tool, input []*message.Message, downstreamAgentOutput []*agent.ResponseUpdate, expectedOutput []*agent.ResponseUpdate, diff --git a/agent/harness/toolautocall/autocall_test.go b/agent/harness/toolautocall/autocall_test.go index e2c7d12a..5d0b8592 100644 --- a/agent/harness/toolautocall/autocall_test.go +++ b/agent/harness/toolautocall/autocall_test.go @@ -39,6 +39,141 @@ func (schemaOnlyTool) ReturnSchema() any { return nil } +func TestFunctionInvoking_InvocationIdentity(t *testing.T) { + toolFailure := errors.New("tool failed") + for _, tc := range []struct { + name string + concurrent bool + wrap bool + additional bool + shadowed bool + emptyID bool + toolError error + }{ + {name: "without wrappers"}, + {name: "middleware wrapper", wrap: true}, + {name: "concurrent wrappers", wrap: true, concurrent: true}, + {name: "empty call ID", wrap: true, emptyID: true}, + {name: "additional tool", additional: true}, + {name: "additional tool wrappers", additional: true, wrap: true}, + {name: "concurrent additional tool wrappers", additional: true, wrap: true, concurrent: true}, + {name: "additional tool failure", additional: true, wrap: true, toolError: toolFailure}, + {name: "request tool takes precedence", shadowed: true, wrap: true}, + {name: "tool failure", wrap: true, toolError: toolFailure}, + } { + t.Run(tc.name, func(t *testing.T) { + type auditKey struct{} + var handlerCalls, wrapperCalls, innerCalls atomic.Int32 + firstID := "call-1" + if tc.emptyID { + firstID = "" + } + testTool := functool.MustNew(functool.Config{Name: "lookup"}, func(ctx context.Context, args struct { + ID string `json:"id"` + }, + ) (string, error) { + handlerCalls.Add(1) + invocation, ok := tool.InvocationFromContext(ctx) + if !ok || invocation.CallID != args.ID { + t.Errorf("handler identity = %#v, %v; want call ID %q", invocation, ok, args.ID) + } + if tc.wrap && ctx.Value(auditKey{}) != args.ID { + t.Errorf("wrapper context value = %v, want %q", ctx.Value(auditKey{}), args.ID) + } + return args.ID, tc.toolError + }) + observer := agent.FunctionInvocationMiddleware(func(ctx context.Context, invocation *agent.FunctionInvocationContext, next agent.FunctionInvocationFunc) (any, error) { + wrapperCalls.Add(1) + identity, ok := tool.InvocationFromContext(ctx) + if !ok || identity.CallID != invocation.CallID || invocation.Function != testTool { + t.Error("callback did not receive the original tool and invocation identity") + } + result, err := next(context.WithValue(ctx, auditKey{}, invocation.CallID), invocation) + if result != invocation.CallID && err == nil { + t.Errorf("wrapper result = %v, want %q", result, invocation.CallID) + } + if !errors.Is(err, tc.toolError) { + t.Errorf("wrapper error = %v, want %v", err, tc.toolError) + } + return result, err + }) + inner := agent.FunctionInvocationMiddleware(func(ctx context.Context, invocation *agent.FunctionInvocationContext, next agent.FunctionInvocationFunc) (any, error) { + innerCalls.Add(1) + if ctx.Value(auditKey{}) != invocation.CallID { + t.Error("callbacks did not execute in registration order") + } + return next(ctx, invocation) + }) + checkProviderContext := func(ctx context.Context, _ []*message.Message, opts ...agent.Option) { + if invocation, ok := tool.InvocationFromContext(ctx); ok { + t.Errorf("invocation leaked to provider: %#v", invocation) + } + if tc.additional && len(slices.Collect(agent.AllOptions(opts, agent.WithTool))) != 0 { + t.Error("additional tools leaked into provider options") + } + } + runner := &agenttest.Runner{Responses: agenttest.NewResponseBuilder(checkProviderContext). + AddFunctionCall(firstID, "lookup", fmt.Sprintf(`{"id":%q}`, firstID)). + AddFunctionCall("call-2", "lookup", `{"id":"call-2"}`). + NewTurn(checkProviderContext).AddText("done").Build()} + cfg := toolautocall.Config{AllowConcurrentInvocations: tc.concurrent} + options := []agent.Option{agent.WithTool(testTool)} + if tc.additional { + cfg.AdditionalTools = []tool.Tool{testTool} + options = nil + } + if tc.shadowed { + cfg.AdditionalTools = []tool.Tool{functool.MustNew(functool.Config{Name: "lookup"}, func(context.Context, struct{}) (string, error) { + t.Error("additional tool took precedence over the request tool") + return "", nil + })} + } + run := toolautocall.New(cfg).Run + next := func(ctx context.Context, messages []*message.Message, options ...agent.Option) iter.Seq2[*agent.ResponseUpdate, error] { + return run(runner.Run, ctx, messages, options...) + } + var updates iter.Seq2[*agent.ResponseUpdate, error] + if tc.wrap { + updates = observer.Run(func(ctx context.Context, messages []*message.Message, opts ...agent.Option) iter.Seq2[*agent.ResponseUpdate, error] { + return inner.Run(next, ctx, messages, opts...) + }, t.Context(), []*message.Message{message.NewText("start")}, options...) + } else { + updates = next(t.Context(), []*message.Message{message.NewText("start")}, options...) + } + var results []*message.FunctionResultContent + for update, err := range updates { + if err != nil { + t.Fatal(err) + } + for _, content := range update.Contents { + if result, ok := content.(*message.FunctionResultContent); ok { + results = append(results, result) + } + } + } + if handlerCalls.Load() != 2 || len(results) != 2 { + t.Fatalf("handler calls = %d, results = %d; want 2 of each", handlerCalls.Load(), len(results)) + } + for i, id := range []string{firstID, "call-2"} { + if results[i].CallID != id || !errors.Is(results[i].Error, tc.toolError) { + t.Errorf("result = %#v, want call ID %q and error %v", results[i], id, tc.toolError) + } + } + if tc.wrap && (wrapperCalls.Load() != 2 || innerCalls.Load() != 2) { + t.Errorf("wrapper calls = %d, inner calls = %d; want 2 of each", wrapperCalls.Load(), innerCalls.Load()) + } + if tc.additional && cfg.AdditionalTools[0] != testTool { + t.Error("middleware mutated configured additional tools") + } + if !tc.additional { + if original, _ := agent.GetOption(options, agent.WithTool); original != testTool { + t.Error("middleware mutated the caller's options") + } + } + }) + } +} + func TestFunctionInvoking_DoesNotInvokeServerHandledFunctionCalls(t *testing.T) { var toolInvoked atomic.Int32 diff --git a/agent/middleware.go b/agent/middleware.go index c1e93837..fce41b6b 100644 --- a/agent/middleware.go +++ b/agent/middleware.go @@ -7,7 +7,9 @@ import ( "iter" "slices" + "github.com/microsoft/agent-framework-go/internal/toolmiddleware" "github.com/microsoft/agent-framework-go/message" + "github.com/microsoft/agent-framework-go/tool" ) // SourceTypeMiddleware represents a message that originated from a middleware component. @@ -38,6 +40,98 @@ func (mf MiddlewareFunc) Run(next RunFunc, ctx context.Context, messages []*mess return mf(next, ctx, messages, options...) } +// FunctionInvocationContext describes a function invocation intercepted by middleware. +type FunctionInvocationContext struct { + // Function is the underlying tool. Treat this field as read-only; changing it + // does not change the tool invoked by the continuation. + Function tool.FuncTool + + // CallID identifies the originating function call. Treat this field as read-only; + // changing it does not change the ID used by the framework for its result. + // It is empty when the provider supplied no ID or there is no invocation context. + CallID string + + // Arguments contains the raw JSON passed to the tool. Middleware may replace + // it before calling next. The tool retains responsibility for input validation. + Arguments string +} + +// FunctionInvocationFunc invokes the next function middleware or underlying tool. +type FunctionInvocationFunc func(context.Context, *FunctionInvocationContext) (any, error) + +// FunctionInvocationMiddleware intercepts function calls by wrapping tools in +// agent options. Call next to continue execution, or return a result/error without +// calling next to replace the tool's behavior. Results and errors may also be +// inspected or replaced after next returns. +// Skipping next replaces only this invocation; it does not terminate the agent loop. +// +// Register it as a [Middleware] before the automatic tool-call middleware and +// after components that supply tools. For tools from context providers, use +// [ProviderConfig.Middlewares]. Function tools configured separately on the +// automatic tool-call middleware are also wrapped, without adding them to provider requests. +// Approval requirements and tool schemas are preserved. Callbacks run only when +// the tool is invoked, including after approval, not when approval is requested. +// +// Multiple callbacks execute in registration order, with the first outermost. +// Each call gets its own FunctionInvocationContext; callbacks must synchronize +// shared application state if tools can execute concurrently. +type FunctionInvocationMiddleware func(ctx context.Context, invocation *FunctionInvocationContext, next FunctionInvocationFunc) (any, error) + +// Run wraps the function tools passed to next without modifying the input options. +func (mf FunctionInvocationMiddleware) Run(next RunFunc, ctx context.Context, messages []*message.Message, options ...Option) iter.Seq2[*ResponseUpdate, error] { + if mf == nil { + return next(ctx, messages, options...) + } + options = slices.Clone(options) + for i, option := range options { + opt, ok := option.(toolOpt) + if !ok { + continue + } + fn, ok := opt.Tool.(tool.FuncTool) + if !ok { + continue + } + options[i] = WithTool(mf.wrap(fn)) + } + options = append(options, toolmiddleware.Wrapper(mf.wrap)) + return next(ctx, messages, options...) +} + +func (mf FunctionInvocationMiddleware) wrap(fn tool.FuncTool) tool.FuncTool { + wrapped := &functionInvocationTool{FuncTool: fn, middlewares: []FunctionInvocationMiddleware{mf}} + if previous, ok := fn.(*functionInvocationTool); ok { + wrapped.FuncTool = previous.FuncTool + wrapped.middlewares = append(slices.Clone(previous.middlewares), mf) + } + return wrapped +} + +type functionInvocationTool struct { + tool.FuncTool + middlewares []FunctionInvocationMiddleware +} + +func (t *functionInvocationTool) ApprovalRequired() bool { + approval, ok := t.FuncTool.(tool.ApprovalRequiredTool) + return ok && approval.ApprovalRequired() +} + +func (t *functionInvocationTool) Call(ctx context.Context, args string) (any, error) { + identity, _ := tool.InvocationFromContext(ctx) + invocation := &FunctionInvocationContext{Function: t.FuncTool, CallID: identity.CallID, Arguments: args} + var next FunctionInvocationFunc = func(ctx context.Context, invocation *FunctionInvocationContext) (any, error) { + return t.FuncTool.Call(ctx, invocation.Arguments) + } + for _, middleware := range slices.Backward(t.middlewares) { + inner := next + next = func(ctx context.Context, invocation *FunctionInvocationContext) (any, error) { + return middleware(ctx, invocation, inner) + } + } + return next(ctx, invocation) +} + // compileRunChain applies the given middlewares around fn. func compileRunChain(fn RunFunc, middlewares []Middleware) RunFunc { for _, mw := range slices.Backward(middlewares) { diff --git a/internal/toolcontext/context.go b/internal/toolcontext/context.go new file mode 100644 index 00000000..b53fa5ac --- /dev/null +++ b/internal/toolcontext/context.go @@ -0,0 +1,19 @@ +// Copyright (c) Microsoft. All rights reserved. + +// Package toolcontext carries tool invocation metadata across framework packages. +package toolcontext + +import "context" + +type callIDKey struct{} + +// WithCallID associates the current tool invocation's call ID with ctx. +func WithCallID(ctx context.Context, callID string) context.Context { + return context.WithValue(ctx, callIDKey{}, callID) +} + +// CallIDFromContext returns the call ID for the current tool invocation. +func CallIDFromContext(ctx context.Context) (string, bool) { + callID, ok := ctx.Value(callIDKey{}).(string) + return callID, ok +} diff --git a/internal/toolmiddleware/wrapper.go b/internal/toolmiddleware/wrapper.go new file mode 100644 index 00000000..9d5b3c5f --- /dev/null +++ b/internal/toolmiddleware/wrapper.go @@ -0,0 +1,13 @@ +// Copyright (c) Microsoft. All rights reserved. + +// Package toolmiddleware shares tool wrappers between agent middleware and tool execution. +package toolmiddleware + +import "github.com/microsoft/agent-framework-go/tool" + +// Wrapper is an internal run option for wrapping tools registered outside agent options. +// Tools already present in agent options are wrapped by the originating middleware. +type Wrapper func(tool.FuncTool) tool.FuncTool + +// MAFValue implements agent.Option without depending on the agent package. +func (w Wrapper) MAFValue() any { return w } diff --git a/tool/tool.go b/tool/tool.go index abd0cab9..1c9318d4 100644 --- a/tool/tool.go +++ b/tool/tool.go @@ -5,6 +5,8 @@ package tool import ( "context" "strings" + + "github.com/microsoft/agent-framework-go/internal/toolcontext" ) // ToolMode represents how tools should be used by the agent. @@ -83,9 +85,33 @@ type FuncTool interface { SchemaTool // Call invokes the tool with raw JSON arguments. + // During automatic tool invocation, [InvocationFromContext] exposes the call ID. Call(ctx context.Context, args string) (any, error) } +// Invocation identifies the logical function call being executed. +// It is a correlation record, not an authorization or cross-run idempotency key. +type Invocation struct { + // CallID is the ID from the originating function call. It matches the ID on + // the function result and is preserved when an approved call resumes. + // An empty provider-supplied ID remains empty; no synthetic ID is assigned. + CallID string +} + +// InvocationFromContext returns the invocation identity supplied by the automatic +// tool-call middleware. Tool wrappers can combine it with their tool reference, +// Call arguments, and returned result or error to correlate lifecycle events. +// The arguments passed to Call remain raw JSON; this accessor does not normalize them. +// +// The identity is scoped to the invocation context, including contexts derived by +// wrappers, and is independent for parallel calls. A direct Call with an unrelated +// context has no invocation identity. Direct calls made with an existing invocation +// context inherit that identity; they do not create a new logical function call. +func InvocationFromContext(ctx context.Context) (Invocation, bool) { + callID, ok := toolcontext.CallIDFromContext(ctx) + return Invocation{CallID: callID}, ok +} + // ApprovalRequiredTool indicates whether a tool requires user approval before invocation. type ApprovalRequiredTool interface { Tool diff --git a/tool/tool_test.go b/tool/tool_test.go index 98ca0d62..21d349ec 100644 --- a/tool/tool_test.go +++ b/tool/tool_test.go @@ -3,11 +3,18 @@ package tool_test import ( + "context" "testing" "github.com/microsoft/agent-framework-go/tool" ) +func TestInvocationFromContext_WithoutInvocation(t *testing.T) { + if invocation, ok := tool.InvocationFromContext(context.Background()); ok || invocation != (tool.Invocation{}) { + t.Fatalf("InvocationFromContext() = %#v, %v; want zero value, false", invocation, ok) + } +} + func TestToolModeRequiredHasNoSpecificTool(t *testing.T) { if name, ok := tool.ToolModeRequired.RequiredTool(); ok { t.Fatalf("RequiredTool() = %q, true; want no specific tool", name)