Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 24 additions & 14 deletions pkg/acp/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -382,6 +382,9 @@ func (a *Agent) runAgent(ctx context.Context, acpSess *Session) error {

eventsChan := acpSess.rt.RunStream(ctx, acpSess.sess)

// Tracks tool call arguments so that we can extract useful information
// once the tool call was made.
toolCallArgs := map[string]string{}
for event := range eventsChan {
if ctx.Err() != nil {
return ctx.Err()
Expand Down Expand Up @@ -411,6 +414,7 @@ func (a *Agent) runAgent(ctx context.Context, acpSess *Session) error {
}

case *runtime.ToolCallEvent:
toolCallArgs[e.ToolCall.ID] = e.ToolCall.Function.Arguments
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Memory leak: toolCallArgs map not cleaned on tool rejection

The toolCallArgs map is populated when a ToolCallEvent is received (line 417), but if the tool call is rejected during the confirmation step (line 410-415), no ToolCallResponseEvent is ever emitted. This means the map entry is never deleted (line 429) and persists for the duration of the agent turn.

While the map is recreated per-turn, this is still a memory leak within long-running turns that process many tool calls with rejections.

Recommendation: Add cleanup in the rejection path:

case *runtime.ToolCallConfirmationEvent:
    resumeType, err := a.handleToolCallConfirmation(ctx, acpSess, e)
    if err != nil {
        return err
    }
    if resumeType == ResumeTypeReject {
        delete(toolCallArgs, e.ToolCall.ID)  // Clean up rejected tool call args
    }

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not true, even when a tool call is rejected we send back a tool call result to the LLM telling it that ... the tool call was rejected.

if err := a.conn.SessionUpdate(ctx, acp.SessionNotification{
SessionId: acp.SessionId(acpSess.id),
Update: buildToolCallStart(e.ToolCall, e.ToolDefinition),
Expand All @@ -419,15 +423,21 @@ func (a *Agent) runAgent(ctx context.Context, acpSess *Session) error {
}

case *runtime.ToolCallResponseEvent:
args, ok := toolCallArgs[e.ToolCallID]
// Should never happen but you know...
if !ok {
return fmt.Errorf("missing tool call arguments for tool call ID %s", e.ToolCallID)
}
delete(toolCallArgs, e.ToolCallID)
if err := a.conn.SessionUpdate(ctx, acp.SessionNotification{
SessionId: acp.SessionId(acpSess.id),
Update: buildToolCallComplete(e.ToolCall, e.Response),
Update: buildToolCallComplete(args, e),
}); err != nil {
return err
}

// Check if this is a todo tool response and emit plan update
if isTodoTool(e.ToolCall.Function.Name) && e.Result != nil && e.Result.Meta != nil {
if isTodoTool(e.ToolDefinition.Name) && e.Result != nil && e.Result.Meta != nil {
if planUpdate := buildPlanUpdateFromTodos(e.Result.Meta); planUpdate != nil {
if err := a.conn.SessionUpdate(ctx, acp.SessionNotification{
SessionId: acp.SessionId(acpSess.id),
Expand Down Expand Up @@ -666,24 +676,24 @@ func determineToolKind(toolName string, tool tools.Tool) acp.ToolKind {
}

// buildToolCallComplete creates a tool call completion update
func buildToolCallComplete(toolCall tools.ToolCall, output string) acp.SessionUpdate {
func buildToolCallComplete(arguments string, event *runtime.ToolCallResponseEvent) acp.SessionUpdate {
// Check if this is a file edit operation and try to extract diff info
if isFileEditTool(toolCall.Function.Name) {
if diffContent := extractDiffContent(toolCall, output); diffContent != nil {
if isFileEditTool(event.ToolDefinition.Name) {
if diffContent := extractDiffContent(event.ToolDefinition.Name, arguments); diffContent != nil {
return acp.UpdateToolCall(
acp.ToolCallId(toolCall.ID),
acp.ToolCallId(event.ToolCallID),
acp.WithUpdateStatus(acp.ToolCallStatusCompleted),
acp.WithUpdateContent([]acp.ToolCallContent{*diffContent}),
acp.WithUpdateRawOutput(map[string]any{"content": output}),
acp.WithUpdateRawOutput(map[string]any{"content": event.Response}),
)
}
}

return acp.UpdateToolCall(
acp.ToolCallId(toolCall.ID),
acp.ToolCallId(event.ToolCallID),
acp.WithUpdateStatus(acp.ToolCallStatusCompleted),
acp.WithUpdateContent([]acp.ToolCallContent{acp.ToolContent(acp.TextBlock(output))}),
acp.WithUpdateRawOutput(map[string]any{"content": output}),
acp.WithUpdateContent([]acp.ToolCallContent{acp.ToolContent(acp.TextBlock(event.Response))}),
acp.WithUpdateRawOutput(map[string]any{"content": event.Response}),
)
}

Expand All @@ -693,8 +703,8 @@ func isFileEditTool(toolName string) bool {
}

// extractDiffContent tries to create a diff content block from edit tool arguments
func extractDiffContent(toolCall tools.ToolCall, _ string) *acp.ToolCallContent {
args := parseToolCallArguments(toolCall.Function.Arguments)
func extractDiffContent(toolCallName, arguments string) *acp.ToolCallContent {
args := parseToolCallArguments(arguments)

// Get the path from arguments
path, ok := args["path"].(string)
Expand All @@ -703,7 +713,7 @@ func extractDiffContent(toolCall tools.ToolCall, _ string) *acp.ToolCallContent
}

// For edit_file, extract the edits
if toolCall.Function.Name == "edit_file" {
if toolCallName == "edit_file" {
edits, ok := args["edits"].([]any)
if !ok || len(edits) == 0 {
return nil
Expand Down Expand Up @@ -735,7 +745,7 @@ func extractDiffContent(toolCall tools.ToolCall, _ string) *acp.ToolCallContent
}

// For write_file, the entire content is new
if toolCall.Function.Name == "write_file" {
if toolCallName == "write_file" {
if content, ok := args["content"].(string); ok {
diff := acp.ToolDiffContent(path, content)
return &diff
Expand Down
4 changes: 2 additions & 2 deletions pkg/cli/printer.go
Original file line number Diff line number Diff line change
Expand Up @@ -134,8 +134,8 @@ func (p *Printer) PrintToolCallWithConfirmation(ctx context.Context, toolCall to
}

// PrintToolCallResponse prints a tool call response
func (p *Printer) PrintToolCallResponse(toolCall tools.ToolCall, response string) {
p.Printf("\n%s response%s\n", bold(toolCall.Function.Name), formatToolCallResponse(response))
func (p *Printer) PrintToolCallResponse(name, response string) {
p.Printf("\n%s response%s\n", bold(name), formatToolCallResponse(response))
}

// PromptMaxIterationsContinue prompts the user to continue after max iterations
Expand Down
4 changes: 2 additions & 2 deletions pkg/cli/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -185,9 +185,9 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess
if cfg.HideToolCalls {
continue
}
out.PrintToolCallResponse(e.ToolCall, e.Response)
out.PrintToolCallResponse(e.ToolDefinition.Name, e.Response)
// Clear the confirmed ID after the tool completes
if e.ToolCall.ID == lastConfirmedToolCallID {
if e.ToolCallID == lastConfirmedToolCallID {
lastConfirmedToolCallID = ""
}
case *runtime.ErrorEvent:
Expand Down
6 changes: 3 additions & 3 deletions pkg/runtime/event.go
Original file line number Diff line number Diff line change
Expand Up @@ -117,19 +117,19 @@ func ToolCallConfirmation(toolCall tools.ToolCall, toolDefinition tools.Tool, ag

type ToolCallResponseEvent struct {
Type string `json:"type"`
ToolCall tools.ToolCall `json:"tool_call"`
ToolCallID string `json:"tool_call_id"`
ToolDefinition tools.Tool `json:"tool_definition"`
Response string `json:"response"`
Result *tools.ToolCallResult `json:"result,omitempty"`
AgentContext
}

func ToolCallResponse(toolCall tools.ToolCall, toolDefinition tools.Tool, result *tools.ToolCallResult, response, agentName string) Event {
func ToolCallResponse(toolCallID string, toolDefinition tools.Tool, result *tools.ToolCallResult, response, agentName string) Event {
return &ToolCallResponseEvent{
Type: "tool_call_response",
ToolCall: toolCall,
Response: response,
Result: result,
ToolCallID: toolCallID,
ToolDefinition: toolDefinition,
AgentContext: newAgentContext(agentName),
}
Expand Down
4 changes: 2 additions & 2 deletions pkg/runtime/tool_dispatch.go
Original file line number Diff line number Diff line change
Expand Up @@ -277,7 +277,7 @@ func (r *LocalRuntime) executeToolWithHandler(
slog.Debug("Tool call completed", "tool", toolCall.Function.Name, "output_length", len(res.Output))
}

events <- ToolCallResponse(toolCall, tool, res, res.Output, a.Name())
events <- ToolCallResponse(toolCall.ID, tool, res, res.Output, a.Name())

// Ensure tool response content is not empty for API compatibility
content := res.Output
Expand Down Expand Up @@ -432,7 +432,7 @@ func addAgentMessage(sess *session.Session, a *agent.Agent, msg *chat.Message, e
// addToolErrorResponse adds a tool error response to the session and emits the event.
// This consolidates the common pattern used by validation, rejection, and cancellation responses.
func (r *LocalRuntime) addToolErrorResponse(_ context.Context, sess *session.Session, toolCall tools.ToolCall, tool tools.Tool, events chan Event, a *agent.Agent, errorMsg string) {
events <- ToolCallResponse(toolCall, tool, tools.ResultError(errorMsg), errorMsg, a.Name())
events <- ToolCallResponse(toolCall.ID, tool, tools.ResultError(errorMsg), errorMsg, a.Name())

toolResponseMsg := chat.Message{
Role: chat.MessageRoleTool,
Expand Down
6 changes: 3 additions & 3 deletions pkg/tui/components/messages/messages.go
Original file line number Diff line number Diff line change
Expand Up @@ -1337,8 +1337,8 @@ func (m *model) AddToolResult(msg *runtime.ToolCallResponseEvent, status types.T
for i := len(m.messages) - 1; i >= 0; i-- {
if m.messages[i].Type == types.MessageTypeAssistantReasoningBlock {
if block, ok := m.views[i].(*reasoningblock.Model); ok {
if block.HasToolCall(msg.ToolCall.ID) {
cmd := block.UpdateToolResult(msg.ToolCall.ID, msg.Response, status, msg.Result)
if block.HasToolCall(msg.ToolCallID) {
cmd := block.UpdateToolResult(msg.ToolCallID, msg.Response, status, msg.Result)
m.invalidateItem(i)
return cmd
}
Expand All @@ -1349,7 +1349,7 @@ func (m *model) AddToolResult(msg *runtime.ToolCallResponseEvent, status types.T
// Then check standalone tool call messages
for i := len(m.messages) - 1; i >= 0; i-- {
toolMessage := m.messages[i]
if toolMessage.Type == types.MessageTypeToolCall && toolMessage.ToolCall.ID == msg.ToolCall.ID {
if toolMessage.Type == types.MessageTypeToolCall && toolMessage.ToolCall.ID == msg.ToolCallID {
toolMessage.Content = strings.ReplaceAll(msg.Response, "\t", " ")
toolMessage.ToolStatus = status
toolMessage.ToolResult = msg.Result
Expand Down
Loading