diff --git a/server/pkg/agent/hermes.go b/server/pkg/agent/hermes.go index c4c9193e42..b25409b556 100644 --- a/server/pkg/agent/hermes.go +++ b/server/pkg/agent/hermes.go @@ -343,6 +343,9 @@ type hermesClient struct { sessionID string onMessage func(Message) onPromptDone func(hermesPromptResult) + // acceptNotification can drop ACP session updates before dispatching to + // handlers that mutate client state such as usage or pending tool calls. + acceptNotification func(updateType string) bool // pendingTools buffers the args for tool calls whose input streams in // across multiple ACP tool_call_update messages (kimi does this — @@ -603,6 +606,9 @@ func (c *hermesClient) handleNotification(raw map[string]json.RawMessage) { } updateType, updateData := normalizeACPUpdate(params.Update) + if c.acceptNotification != nil && !c.acceptNotification(updateType) { + return + } switch updateType { case "agent_message_chunk": diff --git a/server/pkg/agent/kiro.go b/server/pkg/agent/kiro.go index f279f50645..e8e7197182 100644 --- a/server/pkg/agent/kiro.go +++ b/server/pkg/agent/kiro.go @@ -8,6 +8,7 @@ import ( "os/exec" "strings" "sync" + "sync/atomic" "time" ) @@ -83,6 +84,7 @@ func (b *kiroBackend) Execute(ctx context.Context, prompt string, opts ExecOptio var outputMu sync.Mutex var output strings.Builder + var streamingCurrentTurn atomic.Bool promptDone := make(chan hermesPromptResult, 1) @@ -91,7 +93,13 @@ func (b *kiroBackend) Execute(ctx context.Context, prompt string, opts ExecOptio stdin: stdin, pending: make(map[int]*pendingRPC), pendingTools: make(map[string]*pendingToolCall), + acceptNotification: func(string) bool { + return streamingCurrentTurn.Load() + }, onMessage: func(msg Message) { + if !streamingCurrentTurn.Load() { + return + } if msg.Type == MessageToolUse { msg.Tool = kiroToolNameFromTitle(msg.Tool) } @@ -103,6 +111,9 @@ func (b *kiroBackend) Execute(ctx context.Context, prompt string, opts ExecOptio trySend(msgCh, msg) }, onPromptDone: func(result hermesPromptResult) { + if !streamingCurrentTurn.Load() { + return + } select { case promptDone <- result: default: @@ -226,6 +237,7 @@ func (b *kiroBackend) Execute(ctx context.Context, prompt string, opts ExecOptio // requires the standard ACP `prompt` field. Send both so either wire // shape can drive the turn. // TODO: drop one field once Kiro lands on a single canonical payload. + streamingCurrentTurn.Store(true) _, err = c.request(runCtx, "session/prompt", map[string]any{ "sessionId": sessionID, "content": promptBlocks, diff --git a/server/pkg/agent/kiro_test.go b/server/pkg/agent/kiro_test.go index acf7915f44..ec7a77b784 100644 --- a/server/pkg/agent/kiro_test.go +++ b/server/pkg/agent/kiro_test.go @@ -67,6 +67,9 @@ while IFS= read -r line; do printf '{"jsonrpc":"2.0","id":%s,"result":{"sessionId":"ses_new","models":{"currentModelId":"auto","availableModels":[{"modelId":"auto","name":"auto"}]}}}\n' "$id" ;; *'"method":"session/load"'*) + printf '{"jsonrpc":"2.0","method":"session/notification","params":{"sessionId":"ses_loaded","update":{"type":"AgentMessageChunk","content":{"type":"text","text":"history should be ignored"}}}}\n' + printf '{"jsonrpc":"2.0","method":"session/notification","params":{"sessionId":"ses_loaded","update":{"type":"UsageUpdate","usage":{"inputTokens":1000,"outputTokens":1000,"cachedReadTokens":100}}}}\n' + printf '{"jsonrpc":"2.0","method":"session/notification","params":{"sessionId":"ses_loaded","update":{"type":"ToolCall","toolCallId":"tc-current","name":"Shell","status":"pending","parameters":{"command":"echo replay"}}}}\n' printf '{"jsonrpc":"2.0","id":%s,"result":{}}\n' "$id" ;; *'"method":"session/resume"'*) @@ -100,6 +103,7 @@ while IFS= read -r line; do exit 0 ;; esac + printf '{"jsonrpc":"2.0","method":"session/notification","params":{"sessionId":"ses_loaded","update":{"type":"ToolCallUpdate","toolCallId":"tc-current","status":"completed","name":"Shell","parameters":{"command":"echo current"},"output":"current tool output\\n"}}}\n' printf '{"jsonrpc":"2.0","method":"session/notification","params":{"sessionId":"ses_loaded","update":{"type":"AgentMessageChunk","content":{"type":"text","text":"loaded"}}}}\n' printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn","usage":{"inputTokens":2,"outputTokens":1}}}\n' "$id" exit 0 @@ -244,18 +248,47 @@ func TestKiroBackendUsesSessionLoadForResume(t *testing.T) { if err != nil { t.Fatalf("execute: %v", err) } + var messages []Message + messagesDone := make(chan struct{}) go func() { - for range session.Messages { + defer close(messagesDone) + for msg := range session.Messages { + messages = append(messages, msg) } }() result := <-session.Result + <-messagesDone if result.Status != "completed" { t.Fatalf("expected completed result, got status=%q error=%q", result.Status, result.Error) } if result.Output != "loaded" { t.Fatalf("output = %q, want loaded", result.Output) } + if usage := result.Usage["unknown"]; usage.InputTokens != 2 || usage.OutputTokens != 1 || usage.CacheReadTokens != 0 { + t.Fatalf("usage = %+v, want input=2 output=1 cache_read=0", usage) + } + if len(messages) != 3 { + t.Fatalf("messages = %+v, want current tool use, tool result, and text only", messages) + } + if messages[0].Type != MessageToolUse { + t.Fatalf("messages[0].Type = %v, want MessageToolUse", messages[0].Type) + } + if messages[0].Tool != "terminal" { + t.Fatalf("messages[0].Tool = %q, want terminal", messages[0].Tool) + } + if command, _ := messages[0].Input["command"].(string); command != "echo current" { + t.Fatalf("messages[0].Input[command] = %q, want echo current", command) + } + if messages[1].Type != MessageToolResult { + t.Fatalf("messages[1].Type = %v, want MessageToolResult", messages[1].Type) + } + if messages[1].Output != "current tool output\n" { + t.Fatalf("messages[1].Output = %q, want current tool output", messages[1].Output) + } + if messages[2].Type != MessageText || messages[2].Content != "loaded" { + t.Fatalf("messages[2] = %+v, want text loaded", messages[2]) + } if result.SessionID != "ses_existing" { t.Fatalf("session id = %q, want ses_existing", result.SessionID) }