mirror of
https://wget.la/https://github.com/leookun/cursor-byok
synced 2026-08-17 19:47:10 +08:00
Merge pull request #277 from leookun/feat/shell-tool-streaming
feat: add shell tool call delta message handling
This commit is contained in:
@@ -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 统计。
|
// buildTurnEndedMessage 构造 turn 结束消息,并携带标准化后的 token 统计。
|
||||||
func buildTurnEndedMessage(inputTokens int64, outputTokens int64, cacheReadTokens int64, cacheWriteTokens int64) *agentv1.AgentServerMessage {
|
func buildTurnEndedMessage(inputTokens int64, outputTokens int64, cacheReadTokens int64, cacheWriteTokens int64) *agentv1.AgentServerMessage {
|
||||||
inputTokensValue := inputTokens
|
inputTokensValue := inputTokens
|
||||||
|
|||||||
@@ -1003,6 +1003,11 @@ func (service *Service) handleExecResult(intent InboundIntent) error {
|
|||||||
}); err != nil {
|
}); err != nil {
|
||||||
return err
|
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 {
|
if !result.IsTerminal {
|
||||||
return nil
|
return nil
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -145,7 +145,7 @@ var bootstrapStatsigTemplate = statsigBootstrapTemplate{
|
|||||||
bootstrapStatsigGlassCustomThemeSupport: buildEnabledStatsigGate(bootstrapStatsigGlassCustomThemeSupport),
|
bootstrapStatsigGlassCustomThemeSupport: buildEnabledStatsigGate(bootstrapStatsigGlassCustomThemeSupport),
|
||||||
bootstrapStatsigGlassAutomationsUI: buildEnabledStatsigGate(bootstrapStatsigGlassAutomationsUI),
|
bootstrapStatsigGlassAutomationsUI: buildEnabledStatsigGate(bootstrapStatsigGlassAutomationsUI),
|
||||||
bootstrapStatsigTerminalUI2: buildEnabledStatsigGate(bootstrapStatsigTerminalUI2),
|
bootstrapStatsigTerminalUI2: buildEnabledStatsigGate(bootstrapStatsigTerminalUI2),
|
||||||
bootstrapStatsigDisableTerminalOutputUIStreaming: buildEnabledStatsigGate(bootstrapStatsigDisableTerminalOutputUIStreaming),
|
bootstrapStatsigDisableTerminalOutputUIStreaming: buildDisabledStatsigGate(bootstrapStatsigDisableTerminalOutputUIStreaming),
|
||||||
bootstrapStatsigBrowserCanvas: buildEnabledStatsigGate(bootstrapStatsigBrowserCanvas),
|
bootstrapStatsigBrowserCanvas: buildEnabledStatsigGate(bootstrapStatsigBrowserCanvas),
|
||||||
bootstrapStatsigEnableMultitaskMode: buildEnabledStatsigGate(bootstrapStatsigEnableMultitaskMode),
|
bootstrapStatsigEnableMultitaskMode: buildEnabledStatsigGate(bootstrapStatsigEnableMultitaskMode),
|
||||||
bootstrapStatsigDecomposeAlwaysLocalExtHostGate: buildDisabledStatsigGate(bootstrapStatsigDecomposeAlwaysLocalExtHostGate),
|
bootstrapStatsigDecomposeAlwaysLocalExtHostGate: buildDisabledStatsigGate(bootstrapStatsigDecomposeAlwaysLocalExtHostGate),
|
||||||
|
|||||||
@@ -73,3 +73,26 @@ func TestBuildBootstrapStatsigConfigJSONDisablesAlwaysLocalDecompositionGate(t *
|
|||||||
t.Fatalf("unexpected rule_id: %q", ruleID)
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -30,6 +30,7 @@ const (
|
|||||||
var cursorStateDisabledStatsigGates = []string{
|
var cursorStateDisabledStatsigGates = []string{
|
||||||
"decompose_always_local_ext_host",
|
"decompose_always_local_ext_host",
|
||||||
"cursor_extensions_isolation_v2",
|
"cursor_extensions_isolation_v2",
|
||||||
|
"disable_terminal_output_ui_streaming",
|
||||||
}
|
}
|
||||||
|
|
||||||
// InjectCursorUserInfo synchronizes the Cursor user-level auth cache used by the
|
// InjectCursorUserInfo synchronizes the Cursor user-level auth cache used by the
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user