Merge pull request #277 from leookun/feat/shell-tool-streaming

feat: add shell tool call delta message handling

internal/backend/forwarder/events.go

@@ -201,6 +201,41 @@ func buildShellOutputDeltaMessage(delta *agentv1.ShellOutputDeltaUpdate) *agentv
 	}
 }
 
+// buildShellToolCallDeltaMessage maps client shell output to the delta consumed by Cursor's terminal bubble.
+func buildShellToolCallDeltaMessage(callID string, modelCallID string, output *agentv1.ShellOutputDeltaUpdate) *agentv1.AgentServerMessage {
+	if output == nil {
+		return nil
+	}
+	var delta *agentv1.ShellToolCallDelta
+	switch event := output.GetEvent().(type) {
+	case *agentv1.ShellOutputDeltaUpdate_Stdout:
+		content := event.Stdout.GetData()
+		if content == "" {
+			return nil
+		}
+		delta = &agentv1.ShellToolCallDelta{
+			Delta: &agentv1.ShellToolCallDelta_Stdout{
+				Stdout: &agentv1.ShellToolCallStdoutDelta{Content: content},
+			},
+		}
+	case *agentv1.ShellOutputDeltaUpdate_Stderr:
+		content := event.Stderr.GetData()
+		if content == "" {
+			return nil
+		}
+		delta = &agentv1.ShellToolCallDelta{
+			Delta: &agentv1.ShellToolCallDelta_Stderr{
+				Stderr: &agentv1.ShellToolCallStderrDelta{Content: content},
+			},
+		}
+	default:
+		return nil
+	}
+	return buildToolCallDeltaMessage(callID, modelCallID, &agentv1.ToolCallDelta{
+		Delta: &agentv1.ToolCallDelta_ShellToolCallDelta{ShellToolCallDelta: delta},
+	})
+}
+
 // buildTurnEndedMessage 构造 turn 结束消息，并携带标准化后的 token 统计。
 func buildTurnEndedMessage(inputTokens int64, outputTokens int64, cacheReadTokens int64, cacheWriteTokens int64) *agentv1.AgentServerMessage {
 	inputTokensValue := inputTokens

internal/backend/forwarder/service.go

@@ -1003,6 +1003,11 @@ func (service *Service) handleExecResult(intent InboundIntent) error {
 		}); err != nil {
 			return err
 		}
+		if message := buildShellToolCallDeltaMessage(pending.ToolCallID, pending.ModelCallID, result.ShellOutputDelta); message != nil {
+			if err := service.broker.Publish(intent.RequestID, StreamEvent{Message: message}); err != nil {
+				return err
+			}
+		}
 	}
 	if !result.IsTerminal {
 		return nil

internal/backend/forwarder/shell_stream_delta_test.go

@@ -0,0 +1,150 @@
+package forwarder
+
+import (
+	"testing"
+
+	"cursor/gen/agentv1"
+	execbridge "cursor/internal/backend/agent/bridge/exec"
+	runtimecore "cursor/internal/backend/agent/core"
+)
+
+func TestHandleExecResultPublishesShellToolCallDelta(t *testing.T) {
+	tests := []struct {
+		name        string
+		shellStream func() *agentv1.ShellStream
+		wantStdout  string
+		wantStderr  string
+	}{
+		{
+			name: "stdout",
+			shellStream: func() *agentv1.ShellStream {
+				return &agentv1.ShellStream{Event: &agentv1.ShellStream_Stdout{
+					Stdout: &agentv1.ShellStreamStdout{Data: "stdout chunk\n"},
+				}}
+			},
+			wantStdout: "stdout chunk\n",
+		},
+		{
+			name: "stderr",
+			shellStream: func() *agentv1.ShellStream {
+				return &agentv1.ShellStream{Event: &agentv1.ShellStream_Stderr{
+					Stderr: &agentv1.ShellStreamStderr{Data: "stderr chunk\n"},
+				}}
+			},
+			wantStderr: "stderr chunk\n",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			broker := NewStreamBroker()
+			service := &Service{
+				broker:     broker,
+				execBridge: execbridge.NewBridge(),
+			}
+			stream, err := broker.OpenStream(
+				"request-1", "conversation-1", 1, "default", "default",
+				agentv1.AgentMode_AGENT_MODE_AGENT, "run command",
+			)
+			if err != nil {
+				t.Fatalf("OpenStream() error = %v", err)
+			}
+			pending := runtimecore.PendingExec{
+				MessageID:   42,
+				ExecID:      "exec-shell-1",
+				ModelCallID: "model-call-1",
+				ToolCallID:  "tool-call-1",
+				ExecKind:    "shell",
+			}
+			stream.mu.Lock()
+			stream.PendingExecs[pending.ExecID] = pending
+			stream.mu.Unlock()
+
+			if err := service.handleExecResult(InboundIntent{
+				Kind:      "exec_result",
+				RequestID: "request-1",
+				ExecClientMessage: &agentv1.ExecClientMessage{
+					Id:     pending.MessageID,
+					ExecId: pending.ExecID,
+					Message: &agentv1.ExecClientMessage_ShellStream{
+						ShellStream: test.shellStream(),
+					},
+				},
+			}); err != nil {
+				t.Fatalf("handleExecResult() error = %v", err)
+			}
+
+			events, err := broker.ReadFromCursor("request-1", 0)
+			if err != nil {
+				t.Fatalf("ReadFromCursor() error = %v", err)
+			}
+			if len(events) != 2 {
+				t.Fatalf("published events = %d, want compatibility and tool-call deltas", len(events))
+			}
+
+			var compatibilityCount, toolCallDeltaCount int
+			for _, event := range events {
+				update := event.Message.GetInteractionUpdate()
+				if update.GetShellOutputDelta() != nil {
+					compatibilityCount++
+				}
+				deltaUpdate := update.GetToolCallDelta()
+				if deltaUpdate == nil {
+					continue
+				}
+				toolCallDeltaCount++
+				if deltaUpdate.GetCallId() != pending.ToolCallID || deltaUpdate.GetModelCallId() != pending.ModelCallID {
+					t.Fatalf("tool-call delta ids = call %q model %q", deltaUpdate.GetCallId(), deltaUpdate.GetModelCallId())
+				}
+				shellDelta := deltaUpdate.GetToolCallDelta().GetShellToolCallDelta()
+				if shellDelta == nil || shellDelta.GetStdout().GetContent() != test.wantStdout || shellDelta.GetStderr().GetContent() != test.wantStderr {
+					t.Fatalf("shell tool-call delta = %#v", shellDelta)
+				}
+			}
+			if compatibilityCount != 1 || toolCallDeltaCount != 1 {
+				t.Fatalf("published compatibility=%d tool_call_delta=%d, want one each", compatibilityCount, toolCallDeltaCount)
+			}
+		})
+	}
+}
+
+func TestBuildShellToolCallDeltaMessageIgnoresNonOutputEvents(t *testing.T) {
+	tests := []struct {
+		name   string
+		output *agentv1.ShellOutputDeltaUpdate
+	}{
+		{name: "nil"},
+		{
+			name: "start",
+			output: &agentv1.ShellOutputDeltaUpdate{Event: &agentv1.ShellOutputDeltaUpdate_Start{
+				Start: &agentv1.ShellStreamStart{},
+			}},
+		},
+		{
+			name: "exit",
+			output: &agentv1.ShellOutputDeltaUpdate{Event: &agentv1.ShellOutputDeltaUpdate_Exit{
+				Exit: &agentv1.ShellStreamExit{},
+			}},
+		},
+		{
+			name: "empty stdout",
+			output: &agentv1.ShellOutputDeltaUpdate{Event: &agentv1.ShellOutputDeltaUpdate_Stdout{
+				Stdout: &agentv1.ShellStreamStdout{},
+			}},
+		},
+		{
+			name: "empty stderr",
+			output: &agentv1.ShellOutputDeltaUpdate{Event: &agentv1.ShellOutputDeltaUpdate_Stderr{
+				Stderr: &agentv1.ShellStreamStderr{},
+			}},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if message := buildShellToolCallDeltaMessage("tool-call-1", "model-call-1", test.output); message != nil {
+				t.Fatalf("buildShellToolCallDeltaMessage() = %#v, want nil", message)
+			}
+		})
+	}
+}

internal/backend/server/upstream/mocks.go

@@ -145,7 +145,7 @@ var bootstrapStatsigTemplate = statsigBootstrapTemplate{
 		bootstrapStatsigGlassCustomThemeSupport:          buildEnabledStatsigGate(bootstrapStatsigGlassCustomThemeSupport),
 		bootstrapStatsigGlassAutomationsUI:               buildEnabledStatsigGate(bootstrapStatsigGlassAutomationsUI),
 		bootstrapStatsigTerminalUI2:                      buildEnabledStatsigGate(bootstrapStatsigTerminalUI2),
-		bootstrapStatsigDisableTerminalOutputUIStreaming: buildEnabledStatsigGate(bootstrapStatsigDisableTerminalOutputUIStreaming),
+		bootstrapStatsigDisableTerminalOutputUIStreaming: buildDisabledStatsigGate(bootstrapStatsigDisableTerminalOutputUIStreaming),
 		bootstrapStatsigBrowserCanvas:                    buildEnabledStatsigGate(bootstrapStatsigBrowserCanvas),
 		bootstrapStatsigEnableMultitaskMode:              buildEnabledStatsigGate(bootstrapStatsigEnableMultitaskMode),
 		bootstrapStatsigDecomposeAlwaysLocalExtHostGate:  buildDisabledStatsigGate(bootstrapStatsigDecomposeAlwaysLocalExtHostGate),

internal/backend/server/upstream/mocks_test.go

@@ -73,3 +73,26 @@ func TestBuildBootstrapStatsigConfigJSONDisablesAlwaysLocalDecompositionGate(t *
 		t.Fatalf("unexpected rule_id: %q", ruleID)
 	}
 }
+
+func TestBuildBootstrapStatsigConfigJSONEnablesTerminalOutputUIStreaming(t *testing.T) {
+	payload, err := buildBootstrapStatsigConfigJSON(12345, "test-auth-id")
+	if err != nil {
+		t.Fatalf("build bootstrap statsig config: %v", err)
+	}
+
+	var decoded statsigBootstrapTemplate
+	if err := json.Unmarshal(payload, &decoded); err != nil {
+		t.Fatalf("decode bootstrap statsig config: %v", err)
+	}
+
+	gate, ok := decoded.FeatureGates[bootstrapStatsigDisableTerminalOutputUIStreaming]
+	if !ok {
+		t.Fatalf("missing feature gate %q", bootstrapStatsigDisableTerminalOutputUIStreaming)
+	}
+	if value, _ := gate["value"].(bool); value {
+		t.Fatalf("expected %q to be disabled", bootstrapStatsigDisableTerminalOutputUIStreaming)
+	}
+	if ruleID, _ := gate["rule_id"].(string); ruleID != "local_disabled" {
+		t.Fatalf("unexpected rule_id: %q", ruleID)
+	}
+}

internal/cursor/state_db.go

@@ -30,6 +30,7 @@ const (
 var cursorStateDisabledStatsigGates = []string{
 	"decompose_always_local_ext_host",
 	"cursor_extensions_isolation_v2",
+	"disable_terminal_output_ui_streaming",
 }
 
 // InjectCursorUserInfo synchronizes the Cursor user-level auth cache used by the

internal/cursor/state_db_test.go

@@ -0,0 +1,94 @@
+package cursor
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"path/filepath"
+	"testing"
+)
+
+func TestSyncCursorAuthStateDBDisablesCachedTerminalOutputUIStreamingIdempotently(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "state.vscdb")
+	db, err := sql.Open("sqlite", path)
+	if err != nil {
+		t.Fatalf("open temporary state db: %v", err)
+	}
+	if _, err := db.Exec("CREATE TABLE ItemTable (key TEXT UNIQUE ON CONFLICT REPLACE, value BLOB)"); err != nil {
+		db.Close()
+		t.Fatalf("create ItemTable: %v", err)
+	}
+	bootstrap := map[string]any{
+		"feature_gates": map[string]any{
+			"disable_terminal_output_ui_streaming": map[string]any{
+				"value":     true,
+				"rule_id":   "local_enabled",
+				"groupName": "local_enabled",
+			},
+			"unrelated_gate": map[string]any{"value": true},
+		},
+		"hash_used": "none",
+	}
+	raw, err := json.Marshal(bootstrap)
+	if err != nil {
+		db.Close()
+		t.Fatalf("encode bootstrap: %v", err)
+	}
+	if _, err := db.Exec("INSERT INTO ItemTable(key, value) VALUES(?, ?)", cursorStateStatsigBootstrapKey, raw); err != nil {
+		db.Close()
+		t.Fatalf("insert bootstrap: %v", err)
+	}
+	if err := db.Close(); err != nil {
+		t.Fatalf("close setup db: %v", err)
+	}
+
+	values := map[string]string{"cursorAuth/cachedEmail": "local@example.com"}
+	if err := syncCursorAuthStateDB(path, values); err != nil {
+		t.Fatalf("first state sync: %v", err)
+	}
+	first := readCursorStatsigBootstrapForTest(t, path)
+	assertCursorStatsigGateValueForTest(t, first, "disable_terminal_output_ui_streaming", false)
+	assertCursorStatsigGateValueForTest(t, first, "unrelated_gate", true)
+
+	if err := syncCursorAuthStateDB(path, values); err != nil {
+		t.Fatalf("second state sync: %v", err)
+	}
+	second := readCursorStatsigBootstrapForTest(t, path)
+	if string(second) != string(first) {
+		t.Fatalf("repeated sync changed bootstrap:\nfirst:  %s\nsecond: %s", first, second)
+	}
+}
+
+func readCursorStatsigBootstrapForTest(t *testing.T, path string) []byte {
+	t.Helper()
+	db, err := sql.Open("sqlite", path)
+	if err != nil {
+		t.Fatalf("open state db: %v", err)
+	}
+	defer db.Close()
+
+	var raw []byte
+	if err := db.QueryRowContext(context.Background(), "SELECT value FROM ItemTable WHERE key = ?", cursorStateStatsigBootstrapKey).Scan(&raw); err != nil {
+		t.Fatalf("read bootstrap: %v", err)
+	}
+	return raw
+}
+
+func assertCursorStatsigGateValueForTest(t *testing.T, raw []byte, name string, want bool) {
+	t.Helper()
+	var payload struct {
+		FeatureGates map[string]struct {
+			Value bool `json:"value"`
+		} `json:"feature_gates"`
+	}
+	if err := json.Unmarshal(raw, &payload); err != nil {
+		t.Fatalf("decode bootstrap: %v", err)
+	}
+	gate, ok := payload.FeatureGates[name]
+	if !ok {
+		t.Fatalf("missing gate %q", name)
+	}
+	if gate.Value != want {
+		t.Fatalf("gate %q value=%t, want %t", name, gate.Value, want)
+	}
+}