diff --git a/README.md b/README.md index 9b391c7..af26931 100644 --- a/README.md +++ b/README.md @@ -119,14 +119,22 @@ cr init --non-interactive \ ``` Setup with Pi's local RPC runtime. Install Pi's coding agent and make sure the -`pi` binary is available on `PATH` before running `cr review`. New installs -should use the current npm package (`@earendil-works/pi-coding-agent`); existing -installs from the previous npm scope can also work when CR's compatibility -preflight confirms the reviewer controls it requires: RPC/system-prompt mode; -`--no-builtin-tools` with an exact `--tools` allowlist; explicit `--extension` loading while -`--no-extensions` disables discovery; and `--no-context-files`, `--no-approve`, -`--no-skills`, `--no-prompt-templates`, `--no-themes`, and `--no-session`. -CR preflights these capabilities before starting a Pi reviewer and returns an +`pi` binary is available on `PATH` before running `cr review`. Use Pi 1.0 or +newer from the current npm package (`@earendil-works/pi-coding-agent`). CR +relies on Pi 1.0's RPC completion contract: it keeps reading events through +automatic retries and compaction recovery until `agent_settled`, takes the +answer only from the final completed assistant message, and sums usage from +each completed assistant message and compaction summary. A final provider +error or abort fails the task with Pi's diagnostic. Older runtimes, including +installs from the previous npm scope, that end a run without `agent_settled` +fail with an incompatible-runtime error after a short grace period instead of +returning a result. + +CR also preflights the reviewer controls it requires before starting a Pi +reviewer: RPC/system-prompt mode; `--no-builtin-tools` with an exact `--tools` +allowlist; explicit `--extension` loading while `--no-extensions` disables +discovery; and `--no-context-files`, `--no-approve`, `--no-skills`, +`--no-prompt-templates`, `--no-themes`, and `--no-session`. It returns an incompatible-runtime error when any control is unavailable. ```bash diff --git a/internal/llmadapters/pi_rpc.go b/internal/llmadapters/pi_rpc.go index a61d5cd..3d82139 100644 --- a/internal/llmadapters/pi_rpc.go +++ b/internal/llmadapters/pi_rpc.go @@ -14,6 +14,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/open-cli-collective/codereview-cli/internal/llm" @@ -27,9 +28,14 @@ const ( piRPCReviewerToolNames = "cr_read,cr_search,cr_list,cr_diff" piRPCToolEvidenceReserve = 256 piRPCToolDiagnosticMaxRunes = 128 + piRPCErrorMaxRunes = 512 piRPCPreflightTimeout = 5 * time.Second piRPCPreflightOutputBytes = 64 * 1024 piRPCPreflightRegistration = "codereview-pi-reviewer-tools-registered cr_read,cr_search,cr_list,cr_diff" + // piRPCSettlementGrace bounds the quiet period after a run boundary that + // does not announce more work. Pi 1.0 follows it with agent_settled; older + // runtimes never do and would otherwise idle until the task deadline. + piRPCSettlementGrace = 30 * time.Second ) // ErrPiRPCIncompatible reports that the installed Pi runtime cannot enforce @@ -54,6 +60,7 @@ type PiRPCAdapter struct { timeout time.Duration scratchDirFactory ScratchDirFactory fastModeModels []string + settlementGrace time.Duration preflightMu sync.Mutex preflightReady bool } @@ -84,6 +91,7 @@ func NewPiRPCAdapter(opts PiRPCOptions) *PiRPCAdapter { timeout: timeout, scratchDirFactory: factory, fastModeModels: append([]string(nil), opts.FastModeModels...), + settlementGrace: piRPCSettlementGrace, } } @@ -165,6 +173,7 @@ func (a *PiRPCAdapter) Start(ctx context.Context, req Request) (Stream, error) { stdin: process.Stdin(), allowReviewerTools: req.ReviewerWorkspace != nil, logBytesLeft: -1, + settlementGrace: a.settlementGrace, } if req.ReviewerWorkspace != nil { stream.toolEvidenceBytesLeft = min(piRPCToolEvidenceReserve, req.ReviewerWorkspace.MaxToolOutputBytes) @@ -495,6 +504,7 @@ type piRPCStream struct { baseStream stdin io.Closer allowReviewerTools bool + settlementGrace time.Duration logLimitMu sync.Mutex logBytesLeft int logCapped bool @@ -518,10 +528,19 @@ func (s *piRPCStream) run(ctx context.Context, cmd *exec.Cmd, stdout io.Reader, }() scanResult := s.scanStdout(stdout) + // Closing stdin asks Pi to shut down, so it happens only after settlement + // or a terminal failure. Keep reading stdout so shutdown never blocks on a + // full pipe. if s.stdin != nil { _ = s.stdin.Close() } + stdoutDone := make(chan struct{}) + go func() { + defer close(stdoutDone) + _, _ = io.Copy(piRPCLogWriter{stream: s}, stdout) + }() waitErr := cmd.Wait() + <-stdoutDone <-stderrDone evidence := s.reviewerToolEvidence() s.writeReviewerToolEvidence(evidence) @@ -533,18 +552,18 @@ func (s *piRPCStream) run(ctx context.Context, cmd *exec.Cmd, stdout io.Reader, result.err = scanResult.err case errors.Is(ctx.Err(), context.DeadlineExceeded): result.err = ctx.Err() - case scanResult.agentEnd && waitErr != nil: + case scanResult.settled && waitErr != nil: result.err = waitErr - case scanResult.agentEnd && len(scanResult.response.StructuredOutput) == 0: + case scanResult.settled && len(scanResult.response.StructuredOutput) == 0: result.err = errors.New("llm pi rpc: no structured output") - case scanResult.agentEnd: + case scanResult.settled: result.err = nil case ctx.Err() != nil: result.err = ctx.Err() case waitErr != nil: result.err = waitErr - case !scanResult.agentEnd: - result.err = errors.New("llm pi rpc: missing agent_end") + default: + result.err = errors.New("llm pi rpc: stream ended before agent_settled") } recordRequestDuration(&result.response, start, result.err) s.Cancel() @@ -557,14 +576,26 @@ func (s *piRPCStream) run(ctx context.Context, cmd *exec.Cmd, stdout io.Reader, type piRPCScanResult struct { response Response err error - agentEnd bool + settled bool } +// scanStdout consumes RPC records until agent_settled, the only boundary after +// which Pi does no further automatic work. agent_end can be followed by +// retries, compaction recovery, or queued work, and a successful prompt +// response only acknowledges acceptance. func (s *piRPCStream) scanStdout(stdout io.Reader) piRPCScanResult { scanner := bufio.NewScanner(stdout) scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) var result piRPCScanResult + var final *piRPCAssistantMessage + var retryFailure string + settlement := piRPCSettlementWatch{grace: s.settlementGrace, expire: s.Cancel} + defer settlement.stop() for scanner.Scan() { + if settlement.expired.Load() { + break + } + settlement.stop() line := append([]byte(nil), scanner.Bytes()...) s.writeLog(normalizePiRPCLogLine(line)) event, err := parsePiRPCEvent(line) @@ -587,14 +618,30 @@ func (s *piRPCStream) scanStdout(stdout io.Reader) piRPCScanResult { result.err = fmt.Errorf("llm pi rpc: prompt failed: %s", event.responseFailure) return result } - if len(event.structuredOutput) > 0 { - result.response.StructuredOutput = event.structuredOutput + if event.promptHandled { + s.Cancel() + result.err = errors.New("llm pi rpc: prompt was handled without starting a run") + return result + } + result.response.Usage = addPiRPCUsage(result.response.Usage, event.usage) + if event.assistant != nil { + final = event.assistant } - result.response.Usage = mergeUsage(result.response.Usage, event.usage) - if event.agentEnd { - result.agentEnd = true + if event.retryFailure != "" { + retryFailure = event.retryFailure + } + if event.settled { + result.settled = true + result.response.StructuredOutput, result.err = piRPCFinalAnswer(final, retryFailure) return result } + if event.awaitsSettlement { + settlement.arm() + } + } + if settlement.expired.Load() { + result.err = fmt.Errorf("%w: no agent_settled within %s after the run ended; CR requires the Pi 1.0 RPC completion contract", ErrPiRPCIncompatible, s.settlementGrace) + return result } if err := scanner.Err(); err != nil && result.err == nil { result.err = err @@ -602,6 +649,86 @@ func (s *piRPCStream) scanStdout(stdout io.Reader) piRPCScanResult { return result } +// piRPCFinalAnswer accepts only the last completed assistant message, so text +// from an earlier turn or failed attempt never stands in for the answer. +func piRPCFinalAnswer(final *piRPCAssistantMessage, retryFailure string) ([]byte, error) { + switch { + case final == nil: + return nil, nil + case final.stopReason == "error" || final.stopReason == "aborted": + diagnostic := final.errorMessage + if diagnostic == "" { + diagnostic = retryFailure + } + if diagnostic == "" { + diagnostic = "no error message reported" + } + return nil, fmt.Errorf("llm pi rpc: final assistant message %s: %s", final.stopReason, boundPiRPCDiagnostic(diagnostic, piRPCErrorMaxRunes)) + case final.stopReason == "toolUse": + return nil, errors.New("llm pi rpc: run settled without a final answer after tool use") + case final.stopReason == "length": + return nil, errors.New("llm pi rpc: final assistant message was truncated by the length limit") + case final.stopReason == "stop": + return final.output, nil + case final.stopReason == "": + return nil, errors.New("llm pi rpc: final assistant message is missing a stop reason") + default: + return nil, fmt.Errorf("llm pi rpc: unsupported final assistant stop reason %q", final.stopReason) + } +} + +// piRPCSettlementWatch fails a run that goes quiet after a boundary which +// announced no further work but never reports agent_settled. +type piRPCSettlementWatch struct { + grace time.Duration + expire func() + timer *time.Timer + expired atomic.Bool +} + +func (w *piRPCSettlementWatch) arm() { + if w.grace <= 0 { + return + } + w.timer = time.AfterFunc(w.grace, func() { + w.expired.Store(true) + w.expire() + }) +} + +func (w *piRPCSettlementWatch) stop() { + if w.timer != nil { + w.timer.Stop() + w.timer = nil + } +} + +// addPiRPCUsage adds one completed usage report to the run total. Fields that +// no report has provided stay nil (unknown) rather than becoming zero. +func addPiRPCUsage(total Usage, next Usage) Usage { + total.TokensIn = addPiRPCUsageField(total.TokensIn, next.TokensIn) + total.TokensOut = addPiRPCUsageField(total.TokensOut, next.TokensOut) + total.CacheRead = addPiRPCUsageField(total.CacheRead, next.CacheRead) + total.CacheCreate = addPiRPCUsageField(total.CacheCreate, next.CacheCreate) + total.CacheCreate5m = addPiRPCUsageField(total.CacheCreate5m, next.CacheCreate5m) + total.CacheCreate1h = addPiRPCUsageField(total.CacheCreate1h, next.CacheCreate1h) + total.CostUSD = addPiRPCUsageField(total.CostUSD, next.CostUSD) + return total +} + +func addPiRPCUsageField[T int | float64](total *T, next *T) *T { + switch { + case next == nil: + return total + case total == nil: + value := *next + return &value + default: + value := *total + *next + return &value + } +} + type piRPCLogWriter struct{ stream *piRPCStream } func (w piRPCLogWriter) Write(p []byte) (int, error) { @@ -676,7 +803,7 @@ func (s *piRPCStream) reviewerToolEvidence() *llm.ReviewerToolEvidence { } evidence := &llm.ReviewerToolEvidence{DiffStatus: status} if status == llm.DiffToolStatusFailed { - evidence.DiffDiagnostic = boundPiRPCToolError(s.diffToolError) + evidence.DiffDiagnostic = boundPiRPCDiagnostic(s.diffToolError, piRPCToolDiagnosticMaxRunes) } return evidence } @@ -699,13 +826,13 @@ func (s *piRPCStream) writeReviewerToolEvidence(evidence *llm.ReviewerToolEviden } } -func boundPiRPCToolError(value string) string { +func boundPiRPCDiagnostic(value string, maxRunes int) string { value = strings.Join(strings.Fields(strings.TrimSpace(value)), " ") runes := []rune(value) - if len(runes) <= piRPCToolDiagnosticMaxRunes { + if len(runes) <= maxRunes { return value } - return string(runes[:piRPCToolDiagnosticMaxRunes-3]) + "..." + return string(runes[:maxRunes-3]) + "..." } func piRPCToolError(raw json.RawMessage) string { @@ -878,7 +1005,7 @@ func compactPiRPCPartialForLog(partialRaw json.RawMessage) json.RawMessage { type piRPCEvent struct { sessionID string - structuredOutput []byte + assistant *piRPCAssistantMessage usage Usage toolUse bool toolName string @@ -887,7 +1014,18 @@ type piRPCEvent struct { toolFailed bool toolError string responseFailure string - agentEnd bool + promptHandled bool + retryFailure string + awaitsSettlement bool + settled bool +} + +// piRPCAssistantMessage is the authoritative message_end form of one +// completed assistant message. +type piRPCAssistantMessage struct { + output []byte + stopReason string + errorMessage string } func parsePiRPCEvent(line []byte) (piRPCEvent, error) { @@ -902,7 +1040,6 @@ func parsePiRPCEvent(line []byte) (piRPCEvent, error) { eventType := rawString(raw, "type") event := piRPCEvent{ toolUse: piRPCEventIndicatesToolUse(eventType) || valueIndicatesToolUse(decoded), - usage: parsePiRPCUsage(raw), } if event.toolUse { event.toolName = firstRawString(raw, "toolName", "tool_name", "name") @@ -927,26 +1064,41 @@ func parsePiRPCEvent(line []byte) (piRPCEvent, error) { if id := firstRawString(raw, "sessionId", "session_id"); id != "" { event.sessionID = id } - if eventType == "response" && rawString(raw, "command") == "prompt" && !rawBool(raw, "success") { - event.responseFailure = rawString(raw, "error") - if event.responseFailure == "" { - event.responseFailure = "unknown error" + if eventType == "response" && rawString(raw, "command") == "prompt" { + if !rawBool(raw, "success") { + event.responseFailure = rawString(raw, "error") + if event.responseFailure == "" { + event.responseFailure = "unknown error" + } + return event, nil + } + var data map[string]json.RawMessage + if err := json.Unmarshal(raw["data"], &data); err == nil && rawString(data, "disposition") == "handled" { + event.promptHandled = true } return event, nil } - if eventType == "message_end" { - if output, usage := parsePiRPCMessageEnd(raw["message"]); len(output) > 0 { - event.structuredOutput = output - event.usage = mergeUsage(event.usage, usage) - } else { - event.usage = mergeUsage(event.usage, usage) + // Usage is counted only from completed reports: message_end for each + // assistant message and compaction_end for summary generation. + // message_update usage is cumulative, and agent_end.messages repeats + // messages already reported by message_end. + switch eventType { + case "message_end": + event.assistant, event.usage = parsePiRPCMessageEnd(raw["message"]) + case "compaction_end": + var result map[string]json.RawMessage + if err := json.Unmarshal(raw["result"], &result); err == nil { + event.usage = parsePiRPCUsage(result) } - } - if eventType == "agent_end" { - event.agentEnd = true - if output := parsePiRPCAgentEnd(raw["messages"]); len(output) > 0 { - event.structuredOutput = output + event.awaitsSettlement = !rawBool(raw, "willRetry") + case "auto_retry_end": + if !rawBool(raw, "success") { + event.retryFailure = rawString(raw, "finalError") } + case "agent_end": + event.awaitsSettlement = !rawBool(raw, "willRetry") + case "agent_settled": + event.settled = true } return event, nil } @@ -1032,7 +1184,7 @@ func piRPCEventIndicatesToolUse(value string) bool { normalized == "bash" } -func parsePiRPCMessageEnd(value json.RawMessage) ([]byte, Usage) { +func parsePiRPCMessageEnd(value json.RawMessage) (*piRPCAssistantMessage, Usage) { var raw map[string]json.RawMessage if len(value) == 0 { return nil, Usage{} @@ -1040,30 +1192,14 @@ func parsePiRPCMessageEnd(value json.RawMessage) ([]byte, Usage) { if err := json.Unmarshal(value, &raw); err != nil { return nil, Usage{} } - usage := parsePiRPCUsage(raw) if rawString(raw, "role") != "assistant" { - return nil, usage - } - return extractPiRPCText(raw["content"]), usage -} - -func parsePiRPCAgentEnd(value json.RawMessage) []byte { - var messages []map[string]json.RawMessage - if len(value) == 0 { - return nil - } - if err := json.Unmarshal(value, &messages); err != nil { - return nil - } - for i := len(messages) - 1; i >= 0; i-- { - if rawString(messages[i], "role") != "assistant" { - continue - } - if output := extractPiRPCText(messages[i]["content"]); len(output) > 0 { - return output - } + return nil, Usage{} } - return nil + return &piRPCAssistantMessage{ + output: extractPiRPCText(raw["content"]), + stopReason: rawString(raw, "stopReason"), + errorMessage: rawString(raw, "errorMessage"), + }, parsePiRPCUsage(raw) } func extractPiRPCText(value json.RawMessage) []byte { @@ -1099,7 +1235,11 @@ func parsePiRPCUsage(raw map[string]json.RawMessage) Usage { TokensOut: firstRawIntPtr(usageRaw, "tokens_out", "tokensOut", "output", "outputTokens", "completionTokens"), CacheRead: firstRawIntPtr(usageRaw, "cache_read", "cacheRead"), CacheCreate: firstRawIntPtr(usageRaw, "cache_create", "cacheCreate", "cache_write", "cacheWrite"), - CostUSD: firstRawFloatPtr(usageRaw, "cost_usd", "costUSD", "totalCost", "totalCostUSD"), + // Pi reports cacheWrite1h as the one-hour subset of cacheWrite. The + // remainder is not labeled with a retention, so no five-minute value + // is derived from it. + CacheCreate1h: rawIntPtr(usageRaw, "cacheWrite1h"), + CostUSD: firstRawFloatPtr(usageRaw, "cost_usd", "costUSD", "totalCost", "totalCostUSD"), } if usage.CostUSD == nil { usage.CostUSD = nestedRawFloatPtr(usageRaw, "cost", "total") diff --git a/internal/llmadapters/pi_rpc_test.go b/internal/llmadapters/pi_rpc_test.go index c42be9d..c373781 100644 --- a/internal/llmadapters/pi_rpc_test.go +++ b/internal/llmadapters/pi_rpc_test.go @@ -1192,7 +1192,7 @@ func TestPiRPCProtocolFailures(t *testing.T) { } }) - t.Run("caller cancellation propagates before agent_end", func(t *testing.T) { + t.Run("caller cancellation propagates before any event", func(t *testing.T) { recordPath := filepath.Join(t.TempDir(), "record.json") ctx, cancel := context.WithCancel(context.Background()) adapter := NewPiRPCAdapter(PiRPCOptions{ @@ -1233,7 +1233,7 @@ func TestPiRPCProtocolFailures(t *testing.T) { } }) - t.Run("missing agent_end fails stream", func(t *testing.T) { + t.Run("stream ending before agent_settled fails", func(t *testing.T) { recordPath := filepath.Join(t.TempDir(), "record.json") adapter := NewPiRPCAdapter(PiRPCOptions{ Command: os.Args[0], @@ -1246,17 +1246,17 @@ func TestPiRPCProtocolFailures(t *testing.T) { t.Fatalf("Start: %v", err) } _, err = stream.Wait(context.Background()) - if err == nil || !strings.Contains(err.Error(), "agent_end") { - t.Fatalf("Wait error = %v, want missing agent_end", err) + if err == nil || !strings.Contains(err.Error(), "agent_settled") { + t.Fatalf("Wait error = %v, want missing agent_settled", err) } }) - t.Run("non-zero exit after agent_end fails stream", func(t *testing.T) { + t.Run("non-zero exit after agent_settled fails stream", func(t *testing.T) { recordPath := filepath.Join(t.TempDir(), "record.json") adapter := NewPiRPCAdapter(PiRPCOptions{ Command: os.Args[0], commandArgsPrefix: piRPCHelperPrefix(), - Env: piRPCHelperEnv("agent-end-exit-failure", recordPath), + Env: piRPCHelperEnv("settled-exit-failure", recordPath), Timeout: 5 * time.Second, }) stream, err := adapter.Start(context.Background(), Request{Model: "opencode-go/kimi-k2.6", Prompt: "prompt"}) @@ -1269,7 +1269,7 @@ func TestPiRPCProtocolFailures(t *testing.T) { } }) - t.Run("non-zero exit before agent_end fails stream", func(t *testing.T) { + t.Run("non-zero exit before agent_settled fails stream", func(t *testing.T) { recordPath := filepath.Join(t.TempDir(), "record.json") adapter := NewPiRPCAdapter(PiRPCOptions{ Command: os.Args[0], @@ -1288,6 +1288,495 @@ func TestPiRPCProtocolFailures(t *testing.T) { }) } +const piRPCTestPromptStarted = `{"id":"prompt-1","type":"response","command":"prompt","success":true,"data":{"disposition":"started"}}` + +func TestPiRPCSettlementContract(t *testing.T) { + answer := `{"ok":true}` + overloaded := piRPCTestTurn{stopReason: "error", errorMessage: "503 overloaded", usage: piRPCTestUsage(7, 0, nil, 0.25)} + final := piRPCTestTurn{text: answer, stopReason: "stop", usage: piRPCTestUsage(100, 10, map[string]any{"cacheRead": 3, "cacheWrite": 4}, 0.5)} + finalUsage := Usage{TokensIn: piRPCTestPtr(100), TokensOut: piRPCTestPtr(10), CacheRead: piRPCTestPtr(3), CacheCreate: piRPCTestPtr(4), CostUSD: piRPCTestPtr(0.5)} + invalidFinalUsage := Usage{TokensIn: piRPCTestPtr(9), TokensOut: piRPCTestPtr(4), CostUSD: piRPCTestPtr(0.125)} + invalidFinalUsageReport := piRPCTestUsage(9, 4, nil, 0.125) + for _, tt := range []struct { + name string + script []string + wantOutput string + wantErr string + wantUsage Usage + }{ + { + name: "recoverable agent_end continues through retry to settlement", + script: []string{ + piRPCTestPromptStarted, + `{"type":"agent_start"}`, + overloaded.end(), + piRPCTestAgentEnd(true, overloaded), + `#pause`, + `{"type":"auto_retry_start","attempt":1,"maxAttempts":3,"delayMs":100,"errorMessage":"503 overloaded"}`, + `{"type":"agent_start"}`, + `{"type":"message_update","usage":{"input":90,"output":5,"cost":{"total":9}},"assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"{"}}`, + final.end(), + `{"type":"auto_retry_end","success":true,"attempt":1}`, + piRPCTestAgentEnd(false, final), + `{"type":"agent_settled"}`, + }, + wantOutput: answer, + wantUsage: Usage{TokensIn: piRPCTestPtr(107), TokensOut: piRPCTestPtr(10), CacheRead: piRPCTestPtr(3), CacheCreate: piRPCTestPtr(4), CostUSD: piRPCTestPtr(0.75)}, + }, + { + name: "overflow compaction recovery counts summary usage once", + script: []string{ + piRPCTestPromptStarted, + piRPCTestTurn{stopReason: "error", errorMessage: "context overflow", usage: piRPCTestUsage(0, 0, nil, 0)}.end(), + `{"type":"agent_end","messages":[],"willRetry":false}`, + `{"type":"compaction_start","reason":"overflow"}`, + `{"type":"compaction_end","reason":"overflow","result":{"summary":"s","tokensBefore":9000,"usage":{"input":50,"output":5,"cost":{"total":0.25}}},"aborted":false,"willRetry":true}`, + `{"type":"agent_start"}`, + final.end(), + piRPCTestAgentEnd(false, final), + `{"type":"agent_settled"}`, + }, + wantOutput: answer, + wantUsage: Usage{TokensIn: piRPCTestPtr(150), TokensOut: piRPCTestPtr(15), CacheRead: piRPCTestPtr(3), CacheCreate: piRPCTestPtr(4), CostUSD: piRPCTestPtr(0.75)}, + }, + { + name: "completed assistant turns sum usage once and keep only the final answer", + script: []string{ + piRPCTestPromptStarted, + piRPCTestTurn{text: "Checking the diff.", toolCall: true, stopReason: "toolUse", usage: piRPCTestUsage(100, 10, map[string]any{"cacheRead": 1}, 0.25)}.end(), + `{"type":"message_end","message":{"role":"toolResult","toolCallId":"call-1","toolName":"cr_diff","content":[{"type":"text","text":"diff"}],"isError":false}}`, + `{"type":"message_update","usage":{"input":200,"output":20,"cost":{"total":0.5}},"assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"{"}}`, + `{"type":"message_update","usage":{"input":200,"output":20,"cost":{"total":0.5}},"assistantMessageEvent":{"type":"text_end","contentIndex":0,"content":"{\"ok\":true}"}}`, + piRPCTestTurn{text: answer, stopReason: "stop", usage: piRPCTestUsage(200, 20, map[string]any{"cacheRead": 2}, 0.5)}.end(), + `{"type":"agent_end","messages":[{"role":"assistant","content":[],"usage":{"input":100,"output":10},"stopReason":"toolUse"},{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}],"usage":{"input":200,"output":20},"stopReason":"stop"}],"willRetry":false}`, + `{"type":"agent_settled"}`, + }, + wantOutput: answer, + wantUsage: Usage{TokensIn: piRPCTestPtr(300), TokensOut: piRPCTestPtr(30), CacheRead: piRPCTestPtr(3), CostUSD: piRPCTestPtr(0.75)}, + }, + { + name: "one-hour cache writes sum once per completed assistant message", + script: []string{ + piRPCTestPromptStarted, + piRPCTestTurn{text: "Checking the diff.", toolCall: true, stopReason: "toolUse", usage: piRPCTestUsage(100, 10, map[string]any{"cacheWrite": 5, "cacheWrite1h": 1}, 0.25)}.end(), + `{"type":"message_update","usage":{"input":200,"output":20,"cacheWrite":7,"cacheWrite1h":2,"cost":{"total":0.5}},"assistantMessageEvent":{"type":"text_delta","contentIndex":0,"delta":"{"}}`, + piRPCTestTurn{text: answer, stopReason: "stop", usage: piRPCTestUsage(200, 20, map[string]any{"cacheWrite": 7, "cacheWrite1h": 2}, 0.5)}.end(), + `{"type":"agent_end","messages":[{"role":"assistant","content":[],"usage":{"input":100,"output":10,"cacheWrite":5,"cacheWrite1h":1},"stopReason":"toolUse"},{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}],"usage":{"input":200,"output":20,"cacheWrite":7,"cacheWrite1h":2},"stopReason":"stop"}],"willRetry":false}`, + `{"type":"agent_settled"}`, + }, + wantOutput: answer, + wantUsage: Usage{TokensIn: piRPCTestPtr(300), TokensOut: piRPCTestPtr(30), CacheCreate: piRPCTestPtr(12), CacheCreate1h: piRPCTestPtr(3), CostUSD: piRPCTestPtr(0.75)}, + }, + { + name: "one-hour cache split reported by only some messages sums the reported values", + script: []string{ + piRPCTestPromptStarted, + piRPCTestTurn{text: "Checking the diff.", toolCall: true, stopReason: "toolUse", usage: piRPCTestUsage(100, 10, map[string]any{"cacheWrite": 5}, 0.25)}.end(), + piRPCTestTurn{text: answer, stopReason: "stop", usage: piRPCTestUsage(200, 20, map[string]any{"cacheWrite": 7, "cacheWrite1h": 2}, 0.5)}.end(), + piRPCTestAgentEnd(false), + `{"type":"agent_settled"}`, + }, + wantOutput: answer, + wantUsage: Usage{TokensIn: piRPCTestPtr(300), TokensOut: piRPCTestPtr(30), CacheCreate: piRPCTestPtr(12), CacheCreate1h: piRPCTestPtr(2), CostUSD: piRPCTestPtr(0.75)}, + }, + { + name: "provider error after exhausted retries fails with diagnostic", + script: []string{ + piRPCTestPromptStarted, + piRPCTestTurn{text: answer, stopReason: "error", errorMessage: "503 overloaded", usage: piRPCTestUsage(7, 0, nil, 0.25)}.end(), + `{"type":"agent_end","messages":[],"willRetry":true}`, + `{"type":"auto_retry_start","attempt":1,"maxAttempts":1,"delayMs":1,"errorMessage":"503 overloaded"}`, + overloaded.end(), + `{"type":"auto_retry_end","success":false,"attempt":1,"finalError":"503 overloaded"}`, + `{"type":"agent_end","messages":[],"willRetry":false}`, + `{"type":"agent_settled"}`, + }, + wantErr: "503 overloaded", + wantUsage: Usage{TokensIn: piRPCTestPtr(14), TokensOut: piRPCTestPtr(0), CostUSD: piRPCTestPtr(0.5)}, + }, + { + name: "aborted final assistant fails", + script: []string{ + piRPCTestPromptStarted, + piRPCTestTurn{text: `{"ok":`, stopReason: "aborted", errorMessage: "Request was aborted", usage: piRPCTestUsage(5, 1, nil, 0)}.end(), + `{"type":"agent_end","messages":[],"willRetry":false}`, + `{"type":"agent_settled"}`, + }, + wantErr: "Request was aborted", + wantUsage: Usage{TokensIn: piRPCTestPtr(5), TokensOut: piRPCTestPtr(1), CostUSD: piRPCTestPtr(0.0)}, + }, + { + name: "length stop reason rejects truncated final answer", + script: []string{ + piRPCTestPromptStarted, + piRPCTestTurn{text: answer, stopReason: "length", usage: invalidFinalUsageReport}.end(), + piRPCTestAgentEnd(false, piRPCTestTurn{text: answer, stopReason: "length", usage: invalidFinalUsageReport}), + `{"type":"agent_settled"}`, + }, + wantErr: "truncated", + wantUsage: invalidFinalUsage, + }, + { + name: "missing stop reason rejects final answer", + script: []string{ + piRPCTestPromptStarted, + piRPCTestLine(map[string]any{"type": "message_end", "message": map[string]any{ + "role": "assistant", + "content": []map[string]any{{"type": "text", "text": answer}}, + "usage": invalidFinalUsageReport, + }}), + `{"type":"agent_end","messages":[],"willRetry":false}`, + `{"type":"agent_settled"}`, + }, + wantErr: "missing a stop reason", + wantUsage: invalidFinalUsage, + }, + { + name: "empty stop reason rejects final answer", + script: []string{ + piRPCTestPromptStarted, + piRPCTestTurn{text: answer, usage: invalidFinalUsageReport}.end(), + piRPCTestAgentEnd(false, piRPCTestTurn{text: answer, usage: invalidFinalUsageReport}), + `{"type":"agent_settled"}`, + }, + wantErr: "missing a stop reason", + wantUsage: invalidFinalUsage, + }, + { + name: "unknown stop reason rejects final answer", + script: []string{ + piRPCTestPromptStarted, + piRPCTestTurn{text: answer, stopReason: "future_reason", usage: invalidFinalUsageReport}.end(), + piRPCTestAgentEnd(false, piRPCTestTurn{text: answer, stopReason: "future_reason", usage: invalidFinalUsageReport}), + `{"type":"agent_settled"}`, + }, + wantErr: "unsupported final assistant stop reason", + wantUsage: invalidFinalUsage, + }, + { + name: "earlier valid answer then final provider error fails", + script: []string{ + piRPCTestPromptStarted, + piRPCTestTurn{text: answer, toolCall: true, stopReason: "toolUse", usage: piRPCTestUsage(10, 2, nil, 0)}.end(), + piRPCTestTurn{stopReason: "error", errorMessage: "terminated", usage: piRPCTestUsage(0, 0, nil, 0)}.end(), + `{"type":"agent_end","messages":[{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}],"stopReason":"toolUse"}],"willRetry":false}`, + `{"type":"agent_settled"}`, + }, + wantErr: "terminated", + wantUsage: Usage{TokensIn: piRPCTestPtr(10), TokensOut: piRPCTestPtr(2), CostUSD: piRPCTestPtr(0.0)}, + }, + { + name: "settled run without a final answer fails", + script: []string{ + piRPCTestPromptStarted, + piRPCTestTurn{text: "Checking the diff.", toolCall: true, stopReason: "toolUse", usage: piRPCTestUsage(10, 2, nil, 0)}.end(), + `{"type":"agent_end","messages":[],"willRetry":false}`, + `{"type":"agent_settled"}`, + }, + wantErr: "final answer", + wantUsage: Usage{TokensIn: piRPCTestPtr(10), TokensOut: piRPCTestPtr(2), CostUSD: piRPCTestPtr(0.0)}, + }, + { + name: "stream truncated after agent_end fails", + script: []string{ + piRPCTestPromptStarted, + final.end(), + piRPCTestAgentEnd(false, final), + `#exit`, + }, + wantErr: "agent_settled", + wantUsage: finalUsage, + }, + { + name: "malformed record after valid answer fails without output", + script: []string{ + piRPCTestPromptStarted, + final.end(), + `{"type":`, + }, + wantErr: "malformed JSONL", + wantUsage: finalUsage, + }, + { + name: "prompt handled without a run fails", + script: []string{ + `{"id":"prompt-1","type":"response","command":"prompt","success":true,"data":{"disposition":"handled"}}`, + }, + wantErr: "handled", + }, + { + name: "prompt rejection fails", + script: []string{ + `{"id":"prompt-1","type":"response","command":"prompt","success":false,"error":"No API key found for opencode-go"}`, + }, + wantErr: "No API key found", + }, + } { + t.Run(tt.name, func(t *testing.T) { + adapter := piRPCScriptAdapter(t, 5*time.Second, tt.script...) + stream, err := adapter.Start(context.Background(), Request{Model: "opencode-go/kimi-k2.6", Prompt: "prompt"}) + if err != nil { + t.Fatalf("Start: %v", err) + } + start := time.Now() + response, err := stream.Wait(context.Background()) + if elapsed := time.Since(start); elapsed > 3*time.Second { + t.Fatalf("Wait took %s, want completion without waiting for the task deadline", elapsed) + } + if tt.wantErr == "" && err != nil { + t.Fatalf("Wait: %v", err) + } + if tt.wantErr != "" && (err == nil || !strings.Contains(err.Error(), tt.wantErr)) { + t.Fatalf("Wait error = %v, want %q", err, tt.wantErr) + } + if string(response.StructuredOutput) != tt.wantOutput { + t.Fatalf("StructuredOutput = %q, want %q", response.StructuredOutput, tt.wantOutput) + } + if !reflect.DeepEqual(response.Usage, tt.wantUsage) { + t.Fatalf("Usage = %s, want %s", formatPiRPCTestUsage(response.Usage), formatPiRPCTestUsage(tt.wantUsage)) + } + }) + } +} + +func TestParsePiRPCUsageCacheWriteRetention(t *testing.T) { + for _, tt := range []struct { + name string + usage string + want Usage + }{ + { + name: "one-hour subset is kept beside the aggregate", + usage: `{"input":10,"output":2,"cacheRead":3,"cacheWrite":7,"cacheWrite1h":4,"cost":{"total":0.5}}`, + want: Usage{TokensIn: piRPCTestPtr(10), TokensOut: piRPCTestPtr(2), CacheRead: piRPCTestPtr(3), CacheCreate: piRPCTestPtr(7), CacheCreate1h: piRPCTestPtr(4), CostUSD: piRPCTestPtr(0.5)}, + }, + { + name: "zero one-hour subset is a reported zero", + usage: `{"input":10,"output":2,"cacheRead":0,"cacheWrite":7,"cacheWrite1h":0,"cost":{"total":0.5}}`, + want: Usage{TokensIn: piRPCTestPtr(10), TokensOut: piRPCTestPtr(2), CacheRead: piRPCTestPtr(0), CacheCreate: piRPCTestPtr(7), CacheCreate1h: piRPCTestPtr(0), CostUSD: piRPCTestPtr(0.5)}, + }, + { + name: "absent one-hour subset stays unknown", + usage: `{"input":10,"output":2,"cacheRead":0,"cacheWrite":7,"cost":{"total":0.5}}`, + want: Usage{TokensIn: piRPCTestPtr(10), TokensOut: piRPCTestPtr(2), CacheRead: piRPCTestPtr(0), CacheCreate: piRPCTestPtr(7), CostUSD: piRPCTestPtr(0.5)}, + }, + } { + t.Run(tt.name, func(t *testing.T) { + got := parsePiRPCUsage(map[string]json.RawMessage{"usage": json.RawMessage(tt.usage)}) + if !reflect.DeepEqual(got, tt.want) { + t.Fatalf("parsePiRPCUsage = %s, want %s", formatPiRPCTestUsage(got), formatPiRPCTestUsage(tt.want)) + } + }) + } +} + +func TestPiRPCReviewerSettlementPreservesToolEvidence(t *testing.T) { + tempDir := t.TempDir() + repoDir := filepath.Join(tempDir, "repo") + scratchDir := filepath.Join(tempDir, "scratch") + for _, dir := range []string{repoDir, scratchDir} { + if err := os.MkdirAll(dir, 0o700); err != nil { + t.Fatalf("MkdirAll(%s): %v", dir, err) + } + } + diffPath := filepath.Join(tempDir, "diff.patch") + if err := os.WriteFile(diffPath, []byte("fixed diff\n"), 0o600); err != nil { + t.Fatalf("WriteFile(diff): %v", err) + } + commentary := piRPCTestTurn{text: "Checking the diff.", toolCall: true, stopReason: "toolUse", usage: piRPCTestUsage(100, 10, nil, 0.25)} + overloaded := piRPCTestTurn{stopReason: "error", errorMessage: "503 overloaded", usage: piRPCTestUsage(0, 0, nil, 0)} + final := piRPCTestTurn{text: `{"ok":true}`, stopReason: "stop", usage: piRPCTestUsage(200, 20, nil, 0.5)} + adapter := piRPCScriptAdapter(t, 5*time.Second, + piRPCTestPromptStarted, + commentary.end(), + `{"type":"tool_execution_start","toolCallId":"diff-1","toolName":"cr_diff","args":{}}`, + `{"type":"tool_execution_end","toolCallId":"diff-1","toolName":"cr_diff","result":{"content":[{"type":"text","text":"diff"}]},"isError":false}`, + overloaded.end(), + piRPCTestAgentEnd(true, commentary, overloaded), + `#pause`, + `{"type":"auto_retry_start","attempt":1,"maxAttempts":3,"delayMs":100,"errorMessage":"503 overloaded"}`, + final.end(), + `{"type":"auto_retry_end","success":true,"attempt":1}`, + piRPCTestAgentEnd(false, final), + `{"type":"agent_settled"}`, + ) + stream, err := adapter.Start(context.Background(), Request{ + Prompt: "review", + ReviewerWorkspace: &ReviewerWorkspaceRequest{ + RepoDir: repoDir, ScratchDir: scratchDir, DiffPath: diffPath, MaxToolOutputBytes: 2048, + }, + }) + if err != nil { + t.Fatalf("Start: %v", err) + } + response, err := stream.Wait(context.Background()) + if err != nil { + t.Fatalf("Wait: %v", err) + } + if string(response.StructuredOutput) != `{"ok":true}` { + t.Fatalf("StructuredOutput = %q, want recovered final answer", response.StructuredOutput) + } + if response.ReviewerToolEvidence == nil || response.ReviewerToolEvidence.DiffStatus != llm.DiffToolStatusSucceeded { + t.Fatalf("reviewer tool evidence = %#v, want succeeded cr_diff", response.ReviewerToolEvidence) + } + want := Usage{TokensIn: piRPCTestPtr(300), TokensOut: piRPCTestPtr(30), CostUSD: piRPCTestPtr(0.75)} + if !reflect.DeepEqual(response.Usage, want) { + t.Fatalf("Usage = %s, want %s", formatPiRPCTestUsage(response.Usage), formatPiRPCTestUsage(want)) + } +} + +func TestPiRPCPendingRetryHonorsDeadlineAndCancellation(t *testing.T) { + pendingRetry := []string{ + piRPCTestPromptStarted, + piRPCTestTurn{stopReason: "error", errorMessage: "503 overloaded", usage: piRPCTestUsage(0, 0, nil, 0)}.end(), + `{"type":"agent_end","messages":[],"willRetry":true}`, + `{"type":"auto_retry_start","attempt":1,"maxAttempts":3,"delayMs":60000,"errorMessage":"503 overloaded"}`, + } + t.Run("deadline", func(t *testing.T) { + adapter := piRPCScriptAdapter(t, 300*time.Millisecond, pendingRetry...) + stream, err := adapter.Start(context.Background(), Request{Model: "opencode-go/kimi-k2.6", Prompt: "prompt"}) + if err != nil { + t.Fatalf("Start: %v", err) + } + if _, err := stream.Wait(context.Background()); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("Wait error = %v, want deadline exceeded while retry is pending", err) + } + }) + t.Run("cancellation", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + adapter := piRPCScriptAdapter(t, 5*time.Second, pendingRetry...) + stream, err := adapter.Start(ctx, Request{Model: "opencode-go/kimi-k2.6", Prompt: "prompt"}) + if err != nil { + t.Fatalf("Start: %v", err) + } + time.Sleep(300 * time.Millisecond) + cancel() + if _, err := stream.Wait(context.Background()); !errors.Is(err, context.Canceled) { + t.Fatalf("Wait error = %v, want context canceled while retry is pending", err) + } + }) +} + +func TestPiRPCMissingSettlementFailsAsIncompatibleRuntime(t *testing.T) { + final := piRPCTestTurn{text: `{"ok":true}`, stopReason: "stop", usage: piRPCTestUsage(100, 10, nil, 0.5)} + adapter := piRPCScriptAdapter(t, 5*time.Second, + `{"id":"prompt-1","type":"response","command":"prompt","success":true}`, + final.end(), + `{"type":"agent_end","messages":[]}`, + ) + adapter.settlementGrace = 100 * time.Millisecond + stream, err := adapter.Start(context.Background(), Request{Model: "opencode-go/kimi-k2.6", Prompt: "prompt"}) + if err != nil { + t.Fatalf("Start: %v", err) + } + start := time.Now() + response, err := stream.Wait(context.Background()) + if !errors.Is(err, ErrPiRPCIncompatible) || !strings.Contains(err.Error(), "agent_settled") { + t.Fatalf("Wait error = %v, want incompatible runtime without agent_settled", err) + } + if elapsed := time.Since(start); elapsed > 3*time.Second { + t.Fatalf("Wait took %s, want bounded settlement grace", elapsed) + } + if len(response.StructuredOutput) != 0 { + t.Fatalf("StructuredOutput = %q, want none without settlement", response.StructuredOutput) + } +} + +func TestPiRPCSettlementGraceRestartsOnFollowUpWork(t *testing.T) { + final := piRPCTestTurn{text: `{"ok":true}`, stopReason: "stop", usage: piRPCTestUsage(100, 10, nil, 0.5)} + adapter := piRPCScriptAdapter(t, 5*time.Second, + piRPCTestPromptStarted, + final.end(), + piRPCTestAgentEnd(false, final), + `#pause`, + `{"type":"compaction_start","reason":"threshold"}`, + `#pause`, `#pause`, `#pause`, `#pause`, `#pause`, `#pause`, `#pause`, `#pause`, + `{"type":"compaction_end","reason":"threshold","result":{"summary":"s","tokensBefore":9000},"aborted":false,"willRetry":false}`, + `{"type":"agent_settled"}`, + ) + adapter.settlementGrace = 500 * time.Millisecond + stream, err := adapter.Start(context.Background(), Request{Model: "opencode-go/kimi-k2.6", Prompt: "prompt"}) + if err != nil { + t.Fatalf("Start: %v", err) + } + response, err := stream.Wait(context.Background()) + if err != nil { + t.Fatalf("Wait: %v", err) + } + if string(response.StructuredOutput) != `{"ok":true}` { + t.Fatalf("StructuredOutput = %q, want answer after threshold compaction settles", response.StructuredOutput) + } +} + +type piRPCTestTurn struct { + text string + toolCall bool + stopReason string + errorMessage string + usage map[string]any +} + +func (turn piRPCTestTurn) message() map[string]any { + content := []map[string]any{} + if turn.text != "" { + content = append(content, map[string]any{"type": "text", "text": turn.text}) + } + if turn.toolCall { + content = append(content, map[string]any{"type": "toolCall", "id": "call-1", "name": "cr_diff", "arguments": map[string]any{}}) + } + message := map[string]any{"role": "assistant", "content": content, "usage": turn.usage, "stopReason": turn.stopReason} + if turn.errorMessage != "" { + message["errorMessage"] = turn.errorMessage + } + return message +} + +func (turn piRPCTestTurn) end() string { + return piRPCTestLine(map[string]any{"type": "message_end", "message": turn.message()}) +} + +func piRPCTestAgentEnd(willRetry bool, turns ...piRPCTestTurn) string { + messages := make([]map[string]any, 0, len(turns)) + for _, turn := range turns { + messages = append(messages, turn.message()) + } + return piRPCTestLine(map[string]any{"type": "agent_end", "messages": messages, "willRetry": willRetry}) +} + +func piRPCTestUsage(input, output int, cache map[string]any, cost float64) map[string]any { + usage := map[string]any{"input": input, "output": output, "cost": map[string]any{"total": cost}} + for key, value := range cache { + usage[key] = value + } + return usage +} + +func piRPCTestLine(event map[string]any) string { + data, err := json.Marshal(event) + if err != nil { + panic(err) + } + return string(data) +} + +func piRPCTestPtr[T any](value T) *T { return &value } + +func formatPiRPCTestUsage(usage Usage) string { + data, _ := json.Marshal(usage) + return string(data) +} + +func piRPCScriptAdapter(t *testing.T, timeout time.Duration, script ...string) *PiRPCAdapter { + t.Helper() + return NewPiRPCAdapter(PiRPCOptions{ + Command: os.Args[0], + commandArgsPrefix: piRPCHelperPrefix(), + Env: append(piRPCHelperEnv("script", filepath.Join(t.TempDir(), "record.json")), "LLM_PI_RPC_SCRIPT="+strings.Join(script, "\n")), + Timeout: timeout, + }) +} + func TestPiRPCRejectsUnsafeSpecs(t *testing.T) { adapter := NewPiRPCAdapter(PiRPCOptions{}) req := Request{Model: "opencode-go/kimi-k2.6", Prompt: "prompt"} @@ -1417,17 +1906,54 @@ func TestPiRPCHelperProcess(_ *testing.T) { } os.Exit(0) } + // Pi RPC stays alive after agent_settled and shuts down when stdin closes. + stdinClosed := make(chan struct{}) + go func() { + for scanner.Scan() { + } + close(stdinClosed) + }() + awaitStdinClose := func() { + select { + case <-stdinClosed: + case <-time.After(10 * time.Second): + } + } + settle := func() { + fmt.Println(`{"type":"agent_end","messages":[{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}],"stopReason":"stop"}],"willRetry":false}`) + fmt.Println(`{"type":"agent_settled"}`) + awaitStdinClose() + } switch os.Getenv("LLM_HELPER_MODE") { case "success", "slow-success": if os.Getenv("LLM_HELPER_MODE") == "slow-success" { time.Sleep(slowSuccessSleep) } - fmt.Println(`{"id":"prompt-1","type":"response","command":"prompt","success":true}`) + fmt.Println(`{"id":"prompt-1","type":"response","command":"prompt","success":true,"data":{"disposition":"started"}}`) fmt.Println(`{"type":"agent_start","sessionId":"session-1"}`) fmt.Println(`{"type":"message_end","message":{"role":"user","content":"review this diff"}}`) - fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"thinking","text":"ignored"},{"type":"text","text":"{\"ok\":true}"}],"usage":{"tokensIn":3915,"tokensOut":50,"cacheRead":5,"cacheWrite":7,"cost":{"total":0.00392005}}}}`) - fmt.Println(`{"type":"agent_end","messages":[{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}]}`) + fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"thinking","text":"ignored"},{"type":"text","text":"{\"ok\":true}"}],"usage":{"tokensIn":3915,"tokensOut":50,"cacheRead":5,"cacheWrite":7,"cost":{"total":0.00392005}},"stopReason":"stop"}}`) + fmt.Println(`{"type":"agent_end","messages":[{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}],"usage":{"tokensIn":3915,"tokensOut":50,"cacheRead":5,"cacheWrite":7,"cost":{"total":0.00392005}},"stopReason":"stop"}],"willRetry":false}`) + fmt.Println(`{"type":"agent_settled"}`) + awaitStdinClose() + case "script": + for _, line := range strings.Split(os.Getenv("LLM_PI_RPC_SCRIPT"), "\n") { + switch line { + case "#pause": + // Emulate retry delay; Pi shuts down instead of retrying once stdin closes. + select { + case <-stdinClosed: + os.Exit(0) + case <-time.After(100 * time.Millisecond): + } + case "#exit": + os.Exit(0) + default: + fmt.Println(line) + } + } + awaitStdinClose() case "streaming-partials": fmt.Println(`{"id":"prompt-1","type":"response","command":"prompt","success":true}`) fmt.Println(`{"type":"agent_start","sessionId":"session-1"}`) @@ -1472,7 +1998,9 @@ func TestPiRPCHelperProcess(_ *testing.T) { fmt.Println(string(data)) } fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}],"api":"openai-completions","provider":"opencode-go","model":"deepseek-v4-pro","usage":{"input":100,"output":50,"cacheRead":0,"cacheWrite":0,"totalTokens":150,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"stopReason":"stop","timestamp":1700000000000}}`) - fmt.Println(`{"type":"agent_end","messages":[{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}]}`) + fmt.Println(`{"type":"agent_end","messages":[{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}],"willRetry":false}`) + fmt.Println(`{"type":"agent_settled"}`) + awaitStdinClose() case "prompt-failure": fmt.Println(`{"id":"prompt-1","type":"response","command":"prompt","success":false,"error":"No API key found for opencode-go"}`) case "tool": @@ -1485,28 +2013,28 @@ func TestPiRPCHelperProcess(_ *testing.T) { fmt.Printf("{\"type\":\"tool_execution_start\",\"toolCallId\":\"%s\",\"toolName\":%q,\"args\":{}}\n", tool, tool) fmt.Printf("{\"type\":\"tool_execution_end\",\"toolCallId\":\"%s\",\"toolName\":%q,\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"ok\"}]}}\n", tool, tool) } - fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}}`) - fmt.Println(`{"type":"agent_end","messages":[{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}]}`) + fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}],"stopReason":"stop"}}`) + settle() case "reviewer-log-flood": fmt.Fprintln(os.Stderr, strings.Repeat("stderr flood\n", 1000)) fmt.Println(`{"id":"prompt-1","type":"response","command":"prompt","success":true}`) for i := 0; i < 20; i++ { fmt.Printf("{\"type\":\"tool_execution_end\",\"toolCallId\":\"tool-%d\",\"toolName\":\"cr_read\",\"result\":{\"content\":[{\"type\":\"text\",\"text\":%q}]}}\n", i, strings.Repeat("tool output ", 500)) } - fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}}`) - fmt.Println(`{"type":"agent_end","messages":[{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}]}`) + fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}],"stopReason":"stop"}}`) + settle() case "reviewer-diff-failure-log-flood": fmt.Fprintln(os.Stderr, strings.Repeat("stderr flood\n", 1000)) fmt.Println(`{"id":"prompt-1","type":"response","command":"prompt","success":true}`) fmt.Println(`{"type":"tool_execution_start","toolCallId":"diff-1","toolName":"cr_diff","args":{}}`) fmt.Println(`{"type":"tool_execution_end","toolCallId":"diff-1","toolName":"cr_diff","result":{"content":[{"type":"text","text":"fixed diff unavailable"}],"isError":true}}`) - fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}}`) - fmt.Println(`{"type":"agent_end","messages":[{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}]}`) + fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}],"stopReason":"stop"}}`) + settle() case "reviewer-diff-incomplete": fmt.Println(`{"id":"prompt-1","type":"response","command":"prompt","success":true}`) fmt.Println(`{"type":"tool_execution_start","toolCallId":"diff-1","toolName":"cr_diff","args":{}}`) - fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}}`) - fmt.Println(`{"type":"agent_end","messages":[{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}]}`) + fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}],"stopReason":"stop"}}`) + settle() case "sleep": time.Sleep(10 * time.Second) case "malformed": @@ -1514,10 +2042,12 @@ func TestPiRPCHelperProcess(_ *testing.T) { case "no-agent-end": fmt.Println(`{"id":"prompt-1","type":"response","command":"prompt","success":true}`) fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}}`) - case "agent-end-exit-failure": + case "settled-exit-failure": fmt.Println(`{"id":"prompt-1","type":"response","command":"prompt","success":true}`) - fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}}`) - fmt.Println(`{"type":"agent_end","messages":[{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}]}]}`) + fmt.Println(`{"type":"message_end","message":{"role":"assistant","content":[{"type":"text","text":"{\"ok\":true}"}],"stopReason":"stop"}}`) + fmt.Println(`{"type":"agent_end","messages":[],"willRetry":false}`) + fmt.Println(`{"type":"agent_settled"}`) + awaitStdinClose() os.Exit(42) case "exit-before-agent-end": fmt.Println(`{"id":"prompt-1","type":"response","command":"prompt","success":true}`)