From ec70b224d761358bcdfd6a77af10c7e7ef6fca2a Mon Sep 17 00:00:00 2001 From: Kevin Kwok <100260428+Wadehl@users.noreply.github.com> Date: Thu, 24 Sep 2026 16:43:10 +0800 Subject: [PATCH] MUL-7675: feat(grok): support live-run steering via Grok Build ACP (#8796) * feat(grok): support steering active Grok Build runs * test(grok): cover steering before first agent output * fix(grok): bound interject acknowledgement time --- .../internal/handler/task_supplement_test.go | 5 +- server/pkg/agent/grok.go | 97 +++++- server/pkg/agent/grok_test.go | 289 ++++++++++++++++++ server/pkg/agent/hermes.go | 11 + server/pkg/agent/version.go | 4 + server/pkg/agent/version_test.go | 1 + server/pkg/db/generated/supplement.sql.go | 4 +- server/pkg/db/queries/supplement.sql | 4 +- 8 files changed, 407 insertions(+), 8 deletions(-) diff --git a/server/internal/handler/task_supplement_test.go b/server/internal/handler/task_supplement_test.go index 64e6842a0..728192cdc 100644 --- a/server/internal/handler/task_supplement_test.go +++ b/server/internal/handler/task_supplement_test.go @@ -179,6 +179,7 @@ func TestTaskSupplementNegotiationFailsClosed(t *testing.T) { }{ {name: "codex negotiated", provider: "codex", daemonAdvertises: true, wantCapabilityRow: true}, {name: "claude negotiated", provider: "claude", daemonAdvertises: true, wantCapabilityRow: true}, + {name: "grok negotiated", provider: "grok", daemonAdvertises: true, wantCapabilityRow: true}, {name: "claude old daemon", provider: "claude", daemonAdvertises: false, wantCapabilityRow: false}, {name: "old daemon", provider: "codex", daemonAdvertises: false, wantCapabilityRow: false}, {name: "unsupported runtime", provider: "kimi", daemonAdvertises: true, wantCapabilityRow: false}, @@ -211,7 +212,7 @@ func TestTaskSupplementCapabilityDoesNotBreakNonIssueStarts(t *testing.T) { if testHandler == nil { t.Skip("database not available") } - for _, provider := range []string{"codex", "claude"} { + for _, provider := range []string{"codex", "claude", "grok"} { t.Run(provider, func(t *testing.T) { runtimeID := dbfx.Runtime(t, "supplement-no-issue", testutil.Cols{"provider": provider}) agentID := dbfx.Agent(t, "Supplement no issue", runtimeID) @@ -228,7 +229,7 @@ func TestTaskSupplementCapabilityDoesNotBreakNonIssueStarts(t *testing.T) { } func TestTaskSupplementCompletionDoesNotReplayBoundComments(t *testing.T) { - for _, provider := range []string{"codex", "claude"} { + for _, provider := range []string{"codex", "claude", "grok"} { for _, status := range []string{"pending", "delivering", "delivered", "failed"} { for _, ordinary := range []string{"none", "unhandled", "queued"} { t.Run(provider+"/"+status+"/"+ordinary, func(t *testing.T) { diff --git a/server/pkg/agent/grok.go b/server/pkg/agent/grok.go index 9b9438f1c..c3ce9e7bf 100644 --- a/server/pkg/agent/grok.go +++ b/server/pkg/agent/grok.go @@ -2,6 +2,7 @@ package agent import ( "context" + "encoding/json" "fmt" "io" "os/exec" @@ -84,6 +85,83 @@ type grokMessageStream struct { closed bool } +// grokSupplementSession exposes xAI's ACP interjection extension only while +// the matching session/prompt request is active. Grok Build does not advertise +// this vendor extension during initialize, so the daemon separately gates the +// feature by the CLI version that introduced atomic in-turn interjections. +type grokSupplementSession struct { + client *hermesClient + mu sync.RWMutex + sessionID string + active atomic.Bool +} + +// grokSupplementTimeout bounds the wait for an ACP interjection acknowledgement. +var grokSupplementTimeout = 8 * time.Second + +func (s *grokSupplementSession) stop() { + s.active.Store(false) +} + +func (s *grokSupplementSession) ready() bool { + if !s.active.Load() { + return false + } + s.mu.RLock() + defer s.mu.RUnlock() + return s.sessionID != "" +} + +func (s *grokSupplementSession) send(ctx context.Context, text string) error { + if !s.active.Load() { + return fmt.Errorf("grok turn is no longer active") + } + s.mu.RLock() + sessionID := s.sessionID + s.mu.RUnlock() + if sessionID == "" { + return fmt.Errorf("grok session is not ready for interjection") + } + // ACP custom method names are underscore-prefixed on the JSON-RPC wire; + // Grok's ACP server strips that marker before routing to x.ai/interject. + rpcCtx, cancel := context.WithTimeout(ctx, grokSupplementTimeout) + defer cancel() + response, err := s.client.request(rpcCtx, "_x.ai/interject", map[string]any{ + "sessionId": sessionID, + "text": text, + }) + if err != nil { + return fmt.Errorf("grok x.ai/interject failed: %w", err) + } + // Grok's extension wraps its acknowledgement inside the JSON-RPC result: + // {"result":{"status":"queued"}}. The extension can also return an + // in-band error, so a successful JSON-RPC frame alone is not delivery. + var result struct { + Result json.RawMessage `json:"result"` + Error json.RawMessage `json:"error"` + } + if err := json.Unmarshal(response, &result); err != nil { + return fmt.Errorf("grok x.ai/interject returned an invalid response: %w", err) + } + if len(result.Error) > 0 && string(result.Error) != "null" { + var providerError string + if err := json.Unmarshal(result.Error, &providerError); err != nil || providerError == "" { + providerError = "provider extension error" + } + return fmt.Errorf("grok x.ai/interject rejected: %s", providerError) + } + var nested struct { + Status string `json:"status"` + } + if err := json.Unmarshal(result.Result, &nested); err != nil { + return fmt.Errorf("grok x.ai/interject returned an invalid result: %w", err) + } + if nested.Status != "queued" { + return fmt.Errorf("grok x.ai/interject returned status %q, want queued", nested.Status) + } + return nil +} + func newGrokMessageStream(size int) *grokMessageStream { return &grokMessageStream{ch: make(chan Message, size)} } @@ -191,6 +269,7 @@ func (b *grokBackend) Execute(ctx context.Context, prompt string, opts ExecOptio msgStream := newGrokMessageStream(256) resCh := make(chan Result, 1) + supplements := &grokSupplementSession{} // Grok streams interim narration and the final answer as the same // agent_message_chunk type; the tracker keeps only the post-tool-call block @@ -237,6 +316,7 @@ func (b *grokBackend) Execute(ctx context.Context, prompt string, opts ExecOptio } }, } + supplements.client = c readerDone := make(chan struct{}) go func() { @@ -375,6 +455,9 @@ func (b *grokBackend) Execute(ctx context.Context, prompt string, opts ExecOptio } c.sessionID = sessionID + supplements.mu.Lock() + supplements.sessionID = sessionID + supplements.mu.Unlock() b.cfg.Logger.Info("grok session created", "session_id", sessionID) if opts.Model != "" { @@ -425,12 +508,17 @@ func (b *grokBackend) Execute(ctx context.Context, prompt string, opts ExecOptio msgStream.send(Message{Type: MessageStatus, Status: "running", SessionID: sessionID}) streamingCurrentTurn.Store(true) - _, err = c.request(runCtx, "session/prompt", map[string]any{ + _, err = c.requestAndNotifySent(runCtx, "session/prompt", map[string]any{ "sessionId": sessionID, "prompt": []map[string]any{ {"type": "text", "text": userText}, }, + }, func() { + if opts.EnableTaskSupplement { + supplements.active.Store(true) + } }) + supplements.stop() if err != nil { if runCtx.Err() == context.DeadlineExceeded { finalStatus = "timeout" @@ -532,7 +620,12 @@ func (b *grokBackend) Execute(ctx context.Context, prompt string, opts ExecOptio } }() - return &Session{Messages: msgStream.ch, Result: resCh}, nil + session := &Session{Messages: msgStream.ch, Result: resCh} + if opts.EnableTaskSupplement { + session.Supplement = supplements.send + session.SupplementReady = supplements.ready + } + return session, nil } // Grok's ACP `authenticate` method ids (from `initialize`.authMethods). diff --git a/server/pkg/agent/grok_test.go b/server/pkg/agent/grok_test.go index b375ed582..0f50ac5f2 100644 --- a/server/pkg/agent/grok_test.go +++ b/server/pkg/agent/grok_test.go @@ -3,6 +3,7 @@ package agent import ( "context" "encoding/json" + "errors" "log/slog" "os" "path/filepath" @@ -90,6 +91,39 @@ while IFS= read -r line; do esac ;; *'"method":"session/prompt"'*) + if [ -n "$GROK_WAIT_FOR_INTERJECT" ]; then + prompt_id=$id + if [ -z "$GROK_NO_OUTPUT_BEFORE_INTERJECT" ]; then + printf '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"ses_new","update":{"sessionUpdate":"agent_thought_chunk","content":{"type":"text","text":"working"}}}}\n' + fi + while IFS= read -r followup; do + if [ -n "$GROK_REQUESTS_FILE" ]; then + printf '%s\n' "$followup" >> "$GROK_REQUESTS_FILE" + fi + followup_id=$(printf '%s' "$followup" | sed -n 's/.*"id":\([0-9]*\).*/\1/p') + case "$followup" in + *'"method":"_x.ai/interject"'*) + if [ -n "$GROK_INTERJECT_NO_FIRST_RESPONSE" ] && [ -z "$ignored_first_interject" ]; then + ignored_first_interject=1 + continue + fi + if [ -n "$GROK_INTERJECT_UNSUPPORTED" ]; then + printf '{"jsonrpc":"2.0","id":%s,"error":{"code":-32601,"message":"method not found"}}\n' "$followup_id" + elif [ -n "$GROK_INTERJECT_EXTENSION_ERROR" ]; then + printf '{"jsonrpc":"2.0","id":%s,"result":{"result":null,"error":"delivery failed"}}\n' "$followup_id" + else + printf '{"jsonrpc":"2.0","id":%s,"result":{"result":{"status":"%s"}}}\n' "$followup_id" "${GROK_INTERJECT_NESTED_STATUS:-queued}" + fi + if [ -n "$GROK_NO_OUTPUT_BEFORE_INTERJECT" ]; then + printf '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"ses_new","update":{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"111"}}}}\n' + fi + printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn"}}\n' "$prompt_id" + break + ;; + esac + done + exit 0 + fi if [ -n "$GROK_HANG_PROMPT" ]; then while :; do sleep 1; done fi @@ -115,6 +149,261 @@ done ` } +func TestGrokSupplementTargetsActivePrompt(t *testing.T) { + t.Parallel() + fakePath := filepath.Join(t.TempDir(), "grok") + requestsPath := filepath.Join(t.TempDir(), "requests.jsonl") + writeTestExecutable(t, fakePath, []byte(fakeGrokACPScript())) + backend, err := New("grok", Config{ + ExecutablePath: fakePath, + Env: map[string]string{ + "GROK_WAIT_FOR_INTERJECT": "1", + "GROK_NO_OUTPUT_BEFORE_INTERJECT": "1", + "GROK_REQUESTS_FILE": requestsPath, + }, + }) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + session, err := backend.Execute(ctx, "continue working", ExecOptions{Timeout: 5 * time.Second, EnableTaskSupplement: true}) + if err != nil { + t.Fatal(err) + } + if session.Supplement == nil || session.SupplementReady == nil { + t.Fatal("negotiated Grok session did not expose supplement callbacks") + } + if session.SupplementReady() { + t.Fatal("supplement became ready before session/prompt started") + } + deadline := time.Now().Add(3 * time.Second) + for !session.SupplementReady() && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if !session.SupplementReady() { + t.Fatal("supplement never became ready during active prompt") + } + for len(session.Messages) > 0 { + msg := <-session.Messages + if msg.Type != MessageStatus { + t.Fatalf("agent emitted output before interjection: %+v", msg) + } + } + var messages []Message + drained := make(chan struct{}) + go func() { + defer close(drained) + for msg := range session.Messages { + messages = append(messages, msg) + } + }() + if err := session.Supplement(ctx, "Reply with 111 instead of continuing the task."); err != nil { + t.Fatalf("send interjection: %v", err) + } + result := <-session.Result + <-drained + if result.Status != "completed" { + t.Fatalf("turn status=%q error=%q", result.Status, result.Error) + } + if result.Output != "111" { + t.Fatalf("turn output=%q, want steering response 111", result.Output) + } + if len(messages) != 1 || messages[0].Type != MessageText || messages[0].Content != "111" { + t.Fatalf("agent messages=%+v, want only steering response 111", messages) + } + if session.SupplementReady() { + t.Fatal("supplement remained ready after session/prompt completed") + } + requests, err := os.ReadFile(requestsPath) + if err != nil { + t.Fatal(err) + } + for _, want := range []string{`"method":"_x.ai/interject"`, `"sessionId":"ses_new"`, `"text":"Reply with 111 instead of continuing the task."`} { + if !strings.Contains(string(requests), want) { + t.Errorf("ACP requests missing %s:\n%s", want, requests) + } + } +} + +func TestGrokSupplementTimesOutIndependentlyAndAllowsNextMessage(t *testing.T) { + fakePath := filepath.Join(t.TempDir(), "grok") + writeTestExecutable(t, fakePath, []byte(fakeGrokACPScript())) + backend, err := New("grok", Config{ + ExecutablePath: fakePath, + Env: map[string]string{ + "GROK_WAIT_FOR_INTERJECT": "1", + "GROK_INTERJECT_NO_FIRST_RESPONSE": "1", + }, + }) + if err != nil { + t.Fatal(err) + } + + previousTimeout := grokSupplementTimeout + grokSupplementTimeout = 50 * time.Millisecond + defer func() { grokSupplementTimeout = previousTimeout }() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + session, err := backend.Execute(ctx, "continue working", ExecOptions{Timeout: time.Minute, EnableTaskSupplement: true}) + if err != nil { + t.Fatal(err) + } + deadline := time.Now().Add(3 * time.Second) + for !session.SupplementReady() && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if !session.SupplementReady() { + cancel() + t.Fatal("supplement never became ready during active prompt") + } + + started := time.Now() + firstSupplement := make(chan error, 1) + go func() { + firstSupplement <- session.Supplement(ctx, "The first interjection will not be acknowledged.") + }() + select { + case err = <-firstSupplement: + case <-time.After(time.Second): + cancel() + select { + case <-session.Result: + case <-time.After(3 * time.Second): + t.Fatal("Grok run did not stop after its parent context was canceled") + } + t.Fatal("supplement did not return within its independent timeout") + } + if !errors.Is(err, context.DeadlineExceeded) { + cancel() + t.Fatalf("first supplement error=%v, want context deadline exceeded", err) + } + if elapsed := time.Since(started); elapsed > time.Second { + cancel() + t.Fatalf("first supplement took %s to time out, want under 1s", elapsed) + } + if !session.SupplementReady() { + cancel() + t.Fatal("an interject timeout ended the active Grok prompt") + } + + if err := session.Supplement(ctx, "The next interjection should still be delivered."); err != nil { + cancel() + t.Fatalf("second supplement after timeout: %v", err) + } + select { + case result := <-session.Result: + if result.Status != "completed" { + t.Fatalf("turn status=%q error=%q", result.Status, result.Error) + } + case <-time.After(3 * time.Second): + cancel() + t.Fatal("Grok prompt did not complete after the second interjection") + } +} + +func TestGrokSupplementReportsACPRejection(t *testing.T) { + t.Parallel() + fakePath := filepath.Join(t.TempDir(), "grok") + writeTestExecutable(t, fakePath, []byte(fakeGrokACPScript())) + backend, err := New("grok", Config{ + ExecutablePath: fakePath, + Env: map[string]string{ + "GROK_WAIT_FOR_INTERJECT": "1", + "GROK_INTERJECT_UNSUPPORTED": "1", + }, + }) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + session, err := backend.Execute(ctx, "continue working", ExecOptions{Timeout: 5 * time.Second, EnableTaskSupplement: true}) + if err != nil { + t.Fatal(err) + } + drained := make(chan struct{}) + go func() { + defer close(drained) + for range session.Messages { + } + }() + deadline := time.Now().Add(3 * time.Second) + for !session.SupplementReady() && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if !session.SupplementReady() { + t.Fatal("supplement never became ready during active prompt") + } + err = session.Supplement(ctx, "A test message") + if err == nil || !strings.Contains(err.Error(), "method not found") { + t.Fatalf("supplement error=%v, want ACP method-not-found error", err) + } + result := <-session.Result + <-drained + if result.Status != "completed" { + t.Fatalf("turn status=%q error=%q", result.Status, result.Error) + } +} + +func TestGrokSupplementHandlesObservedNestedResultStatus(t *testing.T) { + for _, tc := range []struct { + status string + wantErr string + }{ + {status: "queued"}, + {status: "rejected", wantErr: "rejected"}, + {status: "in_band_error", wantErr: "delivery failed"}, + } { + t.Run(tc.status, func(t *testing.T) { + fakePath := filepath.Join(t.TempDir(), "grok") + writeTestExecutable(t, fakePath, []byte(fakeGrokACPScript())) + env := map[string]string{"GROK_WAIT_FOR_INTERJECT": "1"} + if tc.status == "in_band_error" { + env["GROK_INTERJECT_EXTENSION_ERROR"] = "1" + } else { + env["GROK_INTERJECT_NESTED_STATUS"] = tc.status + } + backend, err := New("grok", Config{ + ExecutablePath: fakePath, + Env: env, + }) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + session, err := backend.Execute(ctx, "continue working", ExecOptions{Timeout: 5 * time.Second, EnableTaskSupplement: true}) + if err != nil { + t.Fatal(err) + } + go func() { + for range session.Messages { + } + }() + deadline := time.Now().Add(3 * time.Second) + for !session.SupplementReady() && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if !session.SupplementReady() { + t.Fatal("supplement never became ready during active prompt") + } + err = session.Supplement(ctx, "A test message") + if tc.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), tc.wantErr) { + t.Fatalf("supplement error=%v, want %q", err, tc.wantErr) + } + } else if err != nil { + t.Fatalf("nested queued status was rejected: %v", err) + } + if result := <-session.Result; result.Status != "completed" { + t.Fatalf("turn status=%q error=%q", result.Status, result.Error) + } + }) + } +} + func TestGrokBackendStreamsAndCompletes(t *testing.T) { t.Parallel() fakePath := filepath.Join(t.TempDir(), "grok") diff --git a/server/pkg/agent/hermes.go b/server/pkg/agent/hermes.go index a9f824804..a5a52179c 100644 --- a/server/pkg/agent/hermes.go +++ b/server/pkg/agent/hermes.go @@ -1048,6 +1048,14 @@ func (c *hermesClient) writeLine(data []byte) error { } func (c *hermesClient) request(ctx context.Context, method string, params any) (json.RawMessage, error) { + return c.requestAndNotifySent(ctx, method, params, nil) +} + +// requestAndNotifySent runs afterWrite once the complete request has been +// written to the provider. Callers that expose a concurrent operation tied to +// that request use it to avoid advertising readiness before the provider has +// received the lifecycle-starting frame. +func (c *hermesClient) requestAndNotifySent(ctx context.Context, method string, params any, afterWrite func()) (json.RawMessage, error) { c.mu.Lock() id := c.nextID c.nextID++ @@ -1075,6 +1083,9 @@ func (c *hermesClient) request(ctx context.Context, method string, params any) ( c.mu.Unlock() return nil, fmt.Errorf("write %s: %w", method, err) } + if afterWrite != nil { + afterWrite() + } select { case res := <-pr.ch: diff --git a/server/pkg/agent/version.go b/server/pkg/agent/version.go index 6cd6a91fc..8fdde27cd 100644 --- a/server/pkg/agent/version.go +++ b/server/pkg/agent/version.go @@ -137,6 +137,10 @@ func SupportsTaskSupplement(provider, version string) bool { // 2.1.110 fixes PreToolUse additionalContext being lost on tool failure. // https://github.com/anthropics/claude-code/blob/main/CHANGELOG.md minimum = "2.1.110" + case "grok": + // Grok Build 1.0.14 supports atomic delivery of in-turn interjections + // through its x.ai/interject ACP extension. + minimum = "1.0.14" default: return false } diff --git a/server/pkg/agent/version_test.go b/server/pkg/agent/version_test.go index f136fb459..dcf48fca5 100644 --- a/server/pkg/agent/version_test.go +++ b/server/pkg/agent/version_test.go @@ -12,6 +12,7 @@ func TestTaskSupplementVersionGate(t *testing.T) { }{ {"codex", "codex-cli 0.100.0", true}, {"codex", "0.99.0", false}, {"claude", "2.1.110 (Claude Code)", true}, {"claude", "2.1.109", false}, + {"grok", "1.0.14", true}, {"grok", "1.0.13", false}, {"grok", "0.2.120", false}, {"codex", "", false}, {"claude", "dev", false}, {"kimi", "9.0.0", false}, } { if got := SupportsTaskSupplement(tc.provider, tc.version); got != tc.want { diff --git a/server/pkg/db/generated/supplement.sql.go b/server/pkg/db/generated/supplement.sql.go index c750f6d21..2bae87485 100644 --- a/server/pkg/db/generated/supplement.sql.go +++ b/server/pkg/db/generated/supplement.sql.go @@ -152,7 +152,7 @@ WITH locked_task AS MATERIALIZED ( AND r.workspace_id = $3 AND t.status = 'running' AND cap.capability = 'task-supplement-v1' - AND r.provider IN ('codex', 'claude') + AND r.provider IN ('codex', 'claude', 'grok') FOR UPDATE OF t ), touched_issue AS ( UPDATE issue i SET @@ -589,7 +589,7 @@ WITH candidate AS MATERIALIZED ( SELECT id, workspace_id, issue_id, 'task-supplement-v1' FROM candidate WHERE $2::boolean - AND provider IN ('codex', 'claude') + AND provider IN ('codex', 'claude', 'grok') AND issue_id IS NOT NULL ON CONFLICT DO NOTHING RETURNING task_id diff --git a/server/pkg/db/queries/supplement.sql b/server/pkg/db/queries/supplement.sql index dea1fb760..b82e4abbf 100644 --- a/server/pkg/db/queries/supplement.sql +++ b/server/pkg/db/queries/supplement.sql @@ -14,7 +14,7 @@ WITH candidate AS MATERIALIZED ( SELECT id, workspace_id, issue_id, 'task-supplement-v1' FROM candidate WHERE @enable_task_supplement::boolean - AND provider IN ('codex', 'claude') + AND provider IN ('codex', 'claude', 'grok') AND issue_id IS NOT NULL ON CONFLICT DO NOTHING RETURNING task_id @@ -45,7 +45,7 @@ WITH locked_task AS MATERIALIZED ( AND r.workspace_id = @workspace_id AND t.status = 'running' AND cap.capability = 'task-supplement-v1' - AND r.provider IN ('codex', 'claude') + AND r.provider IN ('codex', 'claude', 'grok') FOR UPDATE OF t ), touched_issue AS ( UPDATE issue i SET