milestone: 0.7

Change-Id: I3e1e2969b7aa08b53df6b0acb1ef7eda39a6d70b

Author: shentongmartin

adk/agent_tool.go

@@ -91,40 +91,19 @@ func (at *agentTool) Info(ctx context.Context) (*schema.ToolInfo, error) {
 }
 
 func (at *agentTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
-	var intData *agentToolInterruptInfo
-	var bResume bool
-	err := compose.ProcessState(ctx, func(ctx context.Context, s *State) error {
-		toolCallID := compose.GetToolCallID(ctx)
-		intData, bResume = s.AgentToolInterruptData[toolCallID]
-		if bResume {
-			delete(s.AgentToolInterruptData, toolCallID)
-		}
-		return nil
-	})
-	if err != nil {
-		// cannot resume
-		bResume = false
-	}
-
 	var ms *mockStore
 	var iter *AsyncIterator[*AgentEvent]
-	if bResume {
-		ms = newResumeStore(intData.Data)
+	var err error
 
-		iter, err = newInvokableAgentToolRunner(at.agent, ms).Resume(ctx, mockCheckPointID, getOptionsByAgentName(at.agent.Name(ctx), opts)...)
-		if err != nil {
-			return "", err
-		}
-	} else {
+	wasInterrupted, hasState, state := compose.GetInterruptState[[]byte](ctx)
+	if !wasInterrupted {
 		ms = newEmptyStore()
 		var input []Message
 		if at.fullChatHistoryAsInput {
-			history, err := getReactChatHistory(ctx, at.agent.Name(ctx))
+			input, err = getReactChatHistory(ctx, at.agent.Name(ctx))
 			if err != nil {
 				return "", err
 			}
-
-			input = history
 		} else {
 			if at.inputSchema == nil {
 				// default input schema
@@ -144,7 +123,20 @@ func (at *agentTool) InvokableRun(ctx context.Context, argumentsInJSON string, o
 			}
 		}
 
-		iter = newInvokableAgentToolRunner(at.agent, ms).Run(ctx, input, append(getOptionsByAgentName(at.agent.Name(ctx), opts), WithCheckPointID(mockCheckPointID))...)
+		iter = newInvokableAgentToolRunner(at.agent, ms).Run(ctx, input,
+			append(getOptionsByAgentName(at.agent.Name(ctx), opts), WithCheckPointID(mockCheckPointID))...)
+	} else {
+		if !hasState {
+			return "", fmt.Errorf("agent tool '%s' interrupt has happened, but cannot find interrupt state", at.agent.Name(ctx))
+		}
+
+		ms = newResumeStore(state)
+
+		iter, err = newInvokableAgentToolRunner(at.agent, ms).
+			Resume(ctx, mockCheckPointID, getOptionsByAgentName(at.agent.Name(ctx), opts)...)
+		if err != nil {
+			return "", err
+		}
 	}
 
 	var lastEvent *AgentEvent
@@ -169,17 +161,9 @@ func (at *agentTool) InvokableRun(ctx context.Context, argumentsInJSON string, o
 		if !existed {
 			return "", fmt.Errorf("interrupt has happened, but cannot find interrupt info")
 		}
-		err = compose.ProcessState(ctx, func(ctx context.Context, st *State) error {
-			st.AgentToolInterruptData[compose.GetToolCallID(ctx)] = &agentToolInterruptInfo{
-				LastEvent: lastEvent,
-				Data:      data,
-			}
-			return nil
-		})
-		if err != nil {
-			return "", fmt.Errorf("failed to save agent tool checkpoint to state: %w", err)
-		}
-		return "", compose.InterruptAndRerun
+
+		return "", compose.CompositeInterrupt(ctx, "agent tool interrupt", data,
+			lastEvent.Action.internalInterrupted)
 	}
 
 	if lastEvent == nil {

adk/agent_tool_test.go

@@ -217,3 +217,278 @@ func TestGetReactHistory(t *testing.T) {
 		schema.UserMessage("For context: [MyAgent] `transfer_to_agent` tool returned result: successfully transferred to agent [DestAgentName]."),
 	}, result)
 }
+
+// mockAgentWithInputCapture implements the Agent interface for testing and captures the input it receives
+type mockAgentWithInputCapture struct {
+	name          string
+	description   string
+	capturedInput []Message
+	responses     []*AgentEvent
+}
+
+func (a *mockAgentWithInputCapture) Name(_ context.Context) string {
+	return a.name
+}
+
+func (a *mockAgentWithInputCapture) Description(_ context.Context) string {
+	return a.description
+}
+
+func (a *mockAgentWithInputCapture) Run(_ context.Context, input *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {
+	a.capturedInput = input.Messages
+
+	iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
+
+	go func() {
+		defer generator.Close()
+
+		for _, event := range a.responses {
+			generator.Send(event)
+
+			// If the event has an Exit action, stop sending events
+			if event.Action != nil && event.Action.Exit {
+				break
+			}
+		}
+	}()
+
+	return iterator
+}
+
+func newMockAgentWithInputCapture(name, description string, responses []*AgentEvent) *mockAgentWithInputCapture {
+	return &mockAgentWithInputCapture{
+		name:        name,
+		description: description,
+		responses:   responses,
+	}
+}
+
+func TestAgentToolWithOptions(t *testing.T) {
+	// Test Case 1: WithFullChatHistoryAsInput
+	t.Run("WithFullChatHistoryAsInput", func(t *testing.T) {
+		ctx := context.Background()
+
+		// 1. Set up a mock agent that will capture the input it receives
+		mockAgent := newMockAgentWithInputCapture("test-agent", "a test agent", []*AgentEvent{
+			{
+				AgentName: "test-agent",
+				Output: &AgentOutput{
+					MessageOutput: &MessageVariant{
+						IsStreaming: false,
+						Message:     schema.AssistantMessage("done", nil),
+						Role:        schema.Assistant,
+					},
+				},
+			},
+		})
+
+		// 2. Create an agentTool with the option
+		agentTool := NewAgentTool(ctx, mockAgent, WithFullChatHistoryAsInput())
+
+		// 3. Set up a context with a chat history using a graph
+		history := []Message{
+			schema.UserMessage("first user message"),
+			schema.AssistantMessage("first assistant response", nil),
+		}
+
+		g := compose.NewGraph[string, string](compose.WithGenLocalState(func(ctx context.Context) (state *State) {
+			return &State{
+				AgentName: "react-agent",
+				Messages:  append(history, schema.AssistantMessage("tool call msg", nil)),
+			}
+		}))
+
+		assert.NoError(t, g.AddLambdaNode("1", compose.InvokableLambda(func(ctx context.Context, input string) (output string, err error) {
+			// Run the tool within the graph context that has the state
+			_, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{"request":"some ignored input"}`)
+			return "done", err
+		})))
+		assert.NoError(t, g.AddEdge(compose.START, "1"))
+		assert.NoError(t, g.AddEdge("1", compose.END))
+
+		runner, err := g.Compile(ctx)
+		assert.NoError(t, err)
+
+		// 4. Run the graph which will execute the tool with the state
+		_, err = runner.Invoke(ctx, "")
+		assert.NoError(t, err)
+
+		// 5. Assert that the agent received the full history
+		// The agent should receive: history (minus last assistant message) + transfer messages
+		assert.Len(t, mockAgent.capturedInput, 4) // 2 from history + 2 transfer messages
+		assert.Equal(t, "first user message", mockAgent.capturedInput[0].Content)
+		assert.Equal(t, "For context: [react-agent] said: first assistant response.", mockAgent.capturedInput[1].Content)
+		assert.Equal(t, "For context: [react-agent] called tool: `transfer_to_agent` with arguments: test-agent.", mockAgent.capturedInput[2].Content)
+		assert.Equal(t, "For context: [react-agent] `transfer_to_agent` tool returned result: successfully transferred to agent [test-agent].", mockAgent.capturedInput[3].Content)
+	})
+
+	// Test Case 2: WithAgentInputSchema
+	t.Run("WithAgentInputSchema", func(t *testing.T) {
+		ctx := context.Background()
+
+		// 1. Define a custom schema
+		customSchema := schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
+			"custom_arg": {
+				Desc:     "a custom argument",
+				Required: true,
+				Type:     schema.String,
+			},
+		})
+
+		// 2. Set up a mock agent to capture input
+		mockAgent := newMockAgentWithInputCapture("schema-agent", "agent with custom schema", []*AgentEvent{
+			{
+				AgentName: "schema-agent",
+				Output: &AgentOutput{
+					MessageOutput: &MessageVariant{
+						IsStreaming: false,
+						Message:     schema.AssistantMessage("schema processed", nil),
+						Role:        schema.Assistant,
+					},
+				},
+			},
+		})
+
+		// 3. Create agentTool with the custom schema option
+		agentTool := NewAgentTool(ctx, mockAgent, WithAgentInputSchema(customSchema))
+
+		// 4. Verify the Info() method returns the custom schema
+		info, err := agentTool.Info(ctx)
+		assert.NoError(t, err)
+		assert.Equal(t, customSchema, info.ParamsOneOf)
+
+		// 5. Run the tool with arguments matching the custom schema
+		_, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{"custom_arg":"hello world"}`)
+		assert.NoError(t, err)
+
+		// 6. Assert that the agent received the correctly parsed argument
+		// With custom schema, the agent should receive the raw JSON as input
+		assert.Len(t, mockAgent.capturedInput, 1)
+		assert.Equal(t, `{"custom_arg":"hello world"}`, mockAgent.capturedInput[0].Content)
+	})
+
+	// Test Case 3: WithAgentInputSchema with complex schema
+	t.Run("WithAgentInputSchema_ComplexSchema", func(t *testing.T) {
+		ctx := context.Background()
+
+		// 1. Define a complex custom schema with multiple parameters
+		complexSchema := schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
+			"name": {
+				Desc:     "user name",
+				Required: true,
+				Type:     schema.String,
+			},
+			"age": {
+				Desc:     "user age",
+				Required: false,
+				Type:     schema.Integer,
+			},
+			"active": {
+				Desc:     "user status",
+				Required: false,
+				Type:     schema.Boolean,
+			},
+		})
+
+		// 2. Set up a mock agent
+		mockAgent := newMockAgentWithInputCapture("complex-agent", "agent with complex schema", []*AgentEvent{
+			{
+				AgentName: "complex-agent",
+				Output: &AgentOutput{
+					MessageOutput: &MessageVariant{
+						IsStreaming: false,
+						Message:     schema.AssistantMessage("complex processed", nil),
+						Role:        schema.Assistant,
+					},
+				},
+			},
+		})
+
+		// 3. Create agentTool with the complex schema option
+		agentTool := NewAgentTool(ctx, mockAgent, WithAgentInputSchema(complexSchema))
+
+		// 4. Verify the Info() method returns the complex schema
+		info, err := agentTool.Info(ctx)
+		assert.NoError(t, err)
+		assert.Equal(t, complexSchema, info.ParamsOneOf)
+
+		// 5. Run the tool with complex arguments
+		_, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{"name":"John","age":30,"active":true}`)
+		assert.NoError(t, err)
+
+		// 6. Assert that the agent received the complex JSON
+		assert.Len(t, mockAgent.capturedInput, 1)
+		assert.Equal(t, `{"name":"John","age":30,"active":true}`, mockAgent.capturedInput[0].Content)
+	})
+
+	// Test Case 4: Both options together
+	t.Run("BothOptionsTogether", func(t *testing.T) {
+		ctx := context.Background()
+
+		// 1. Define a custom schema
+		customSchema := schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
+			"query": {
+				Desc:     "search query",
+				Required: true,
+				Type:     schema.String,
+			},
+		})
+
+		// 2. Set up a mock agent
+		mockAgent := newMockAgentWithInputCapture("combined-agent", "agent with both options", []*AgentEvent{
+			{
+				AgentName: "combined-agent",
+				Output: &AgentOutput{
+					MessageOutput: &MessageVariant{
+						IsStreaming: false,
+						Message:     schema.AssistantMessage("combined processed", nil),
+						Role:        schema.Assistant,
+					},
+				},
+			},
+		})
+
+		// 3. Create agentTool with both options
+		agentTool := NewAgentTool(ctx, mockAgent, WithAgentInputSchema(customSchema), WithFullChatHistoryAsInput())
+
+		// 4. Set up a context with chat history using a graph
+		history := []Message{
+			schema.UserMessage("previous conversation"),
+			schema.AssistantMessage("previous response", nil),
+		}
+
+		g := compose.NewGraph[string, string](compose.WithGenLocalState(func(ctx context.Context) (state *State) {
+			return &State{
+				AgentName: "react-agent",
+				Messages:  append(history, schema.AssistantMessage("tool call", nil)),
+			}
+		}))
+
+		assert.NoError(t, g.AddLambdaNode("1", compose.InvokableLambda(func(ctx context.Context, input string) (output string, err error) {
+			// Run the tool within the graph context that has the state
+			_, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{"query":"current query"}`)
+			return "done", err
+		})))
+		assert.NoError(t, g.AddEdge(compose.START, "1"))
+		assert.NoError(t, g.AddEdge("1", compose.END))
+
+		runner, err := g.Compile(ctx)
+		assert.NoError(t, err)
+
+		// 5. Run the graph which will execute the tool with the state
+		_, err = runner.Invoke(ctx, "")
+		assert.NoError(t, err)
+
+		// 6. Verify both options work together
+		info, err := agentTool.Info(ctx)
+		assert.NoError(t, err)
+		assert.Equal(t, customSchema, info.ParamsOneOf)
+
+		// The agent should receive full history + the custom query
+		assert.Len(t, mockAgent.capturedInput, 4) // 2 history + 2 transfer messages
+		assert.Equal(t, "previous conversation", mockAgent.capturedInput[0].Content)
+		assert.Equal(t, "For context: [react-agent] said: previous response.", mockAgent.capturedInput[1].Content)
+		assert.Equal(t, "For context: [react-agent] called tool: `transfer_to_agent` with arguments: combined-agent.", mockAgent.capturedInput[2].Content)
+		assert.Equal(t, "For context: [react-agent] `transfer_to_agent` tool returned result: successfully transferred to agent [combined-agent].", mockAgent.capturedInput[3].Content)
+	})
+}