Add function invocation middleware with call identity

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
 }

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()

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,