package agent import ( "context" "os" "path/filepath" "testing" "time" "github.com/stretchr/testify/require" "github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract" "github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig" "github.com/tidbcloud/AgentTeam/desktop/internal/consts" ) func TestPiRPCSpawnWaitsForInitialCorrelatedAck(t *testing.T) { marker := "" usage := piFakeUsage(1, 1, 0, 2) scenario := []piFakeScript{ piFakePromptScript( piFakeUserMessage("initial-ack", "first\n\n"+marker), piFakeUserMessageEnd("initial-ack", "first\n\n"+marker), piFakeAssistantEnd([]map[string]any{piFakeTextBlock("ok")}, usage), piFakeTurnEnd([]map[string]any{piFakeTextBlock("ok")}, usage), piFakeAgentEnd(false), piFakeAgentSettled(), ), } binary, _ := writeFakePiRPCModeBinary(t, scenario) cfg := piTestConfig(t, binary) setTestLaunchSurface(&cfg, "Prompt", "first") cfg.InitialDeliveryMessageID = "initial-ack" result, err := (&PiRPCDriver{}).Spawn(t.Context(), cfg) require.NoError(t, err) t.Cleanup(result.Cleanup) require.Equal(t, "initial-ack", cfg.InitialDeliveryMessageID) } func TestPiRPCSpawnRejectsInitialPromptThatSettlesWithoutAck(t *testing.T) { scenario := []piFakeScript{ piFakePromptScript( piFakeAgentStart(), piFakeAgentEnd(false), piFakeAgentSettled(), ), } binary, _ := writeFakePiRPCModeBinary(t, scenario) cfg := piTestConfig(t, binary) setTestLaunchSurface(&cfg, "Prompt", "first") cfg.InitialDeliveryMessageID = "initial-never-acked" _, err := (&PiRPCDriver{}).Spawn(t.Context(), cfg) require.Error(t, err) require.Contains(t, err.Error(), "settled before runtime input acknowledgement") } func TestPiRPCPlainStringAckCanArriveAfterRootSettlement(t *testing.T) { marker := "late input\n\n" scenario := []piFakeScript{ piFakePromptScript( piFakeAgentStart(), piFakeAgentEnd(false), piFakeAgentSettled(), // Pi normally writes this user frame before assistant work, but a // delayed JSONL flush must remain an authoritative ack even after // the root terminal boundary. piFakePlainUserMessage(marker), piFakePlainUserMessageEnd(marker), ), } binary, _ := writeFakePiRPCModeBinary(t, scenario) d := &PiRPCDriver{} result, err := d.Spawn(t.Context(), piTestConfig(t, binary)) require.NoError(t, err) t.Cleanup(result.Cleanup) require.NoError(t, d.DeliverRuntimeInput(t.Context(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "plain-late"}, "late input"), result.SessionID)) events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeInputAcked) count := 0 for _, event := range events { if event.kind == consts.ParsedEventKindRuntimeInputAcked && event.runtimeInputID == "plain-late" { count++ } } require.Equal(t, 1, count, "plain-string user content is acknowledged exactly once after settlement") } func TestPiRPCReconcileRuntimeInputUsesDurableSessionMarker(t *testing.T) { stateDir := t.TempDir() cfg := drivercontract.LaunchConfig{ RuntimeType: runtimeconfig.TypePi, RuntimeStateDir: stateDir, WorkDir: t.TempDir(), } sessionDir := filepath.Join(stateDir, ".pi-sessions") require.NoError(t, os.MkdirAll(sessionDir, 0o700)) sessionPath := filepath.Join(sessionDir, "2026_session-reconcile.jsonl") require.NoError(t, os.WriteFile(sessionPath, []byte(`{"type":"message","message":{"role":"user","content":[{"type":"text","text":"prompt\n\n"}]}}`+"\n"), 0o600)) driver := &PiRPCDriver{} driver.launchCfg = cfg driver.sessionID = "session-reconcile" got, err := driver.ReconcileRuntimeInput(t.Context(), "seen", driver.sessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationAcked, got) got, err = driver.ReconcileRuntimeInput(t.Context(), "absent", driver.sessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationUnknown, got, "readable EOF without an input-specific negative fence is not authoritative absent") // A replacement can be live and quiescent before the old process has // flushed the first JSONL file. Missing/non-materialized session evidence // is therefore unknown, never authoritative absent (and never replayable). require.NoError(t, os.Remove(sessionPath)) got, err = driver.ReconcileRuntimeInput(t.Context(), "missing-session", driver.sessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationUnknown, got) require.NoError(t, os.WriteFile(sessionPath, []byte(`{"type":"message","message":{"role":"assistant","content":[{"type":"text","text":"quoted\n\n"}]}}`+"\n"), 0o600)) got, err = driver.ReconcileRuntimeInput(t.Context(), "assistant-only", driver.sessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationUnknown, got, "assistant-only content is not input-specific rejection evidence") corruptJSONL := `{"type":"message","message":{"role":"user","content":"prompt\n\n"}}` + "\n" + `{"type":"message"}` + "\n" + `{"type":"message",` require.NoError(t, os.WriteFile(sessionPath, []byte(corruptJSONL), 0o600)) got, err = driver.ReconcileRuntimeInput(t.Context(), "corrupt", driver.sessionID) require.Error(t, err) require.Equal(t, RuntimeInputReconciliationUnknown, got, "a positive marker in a damaged JSONL snapshot is not authoritative") require.NoError(t, os.WriteFile(sessionPath, []byte(`{"type":"session","id":"session-reconcile"}`+"\n"), 0o600)) driver.mu.Lock() driver.rootRunActive = true driver.mu.Unlock() got, err = driver.ReconcileRuntimeInput(t.Context(), "pending", driver.sessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationPending, got) driver.mu.Lock() driver.rootRunActive = false driver.transportClosed = true driver.mu.Unlock() got, err = driver.ReconcileRuntimeInput(t.Context(), "closed", driver.sessionID) require.ErrorIs(t, err, errPiRPCTransportClosed) require.Equal(t, RuntimeInputReconciliationUnknown, got) } func TestPiRPCReconcileReadableEOFIsFailClosedUntilMarkerMaterializes(t *testing.T) { stateDir := t.TempDir() cfg := drivercontract.LaunchConfig{ RuntimeType: runtimeconfig.TypePi, RuntimeStateDir: stateDir, WorkDir: t.TempDir(), } sessionDir := filepath.Join(stateDir, ".pi-sessions") require.NoError(t, os.MkdirAll(sessionDir, 0o700)) sessionPath := filepath.Join(sessionDir, "2026_session-appendable.jsonl") tests := []struct { name string line string }{ { name: "header-only", line: `{"type":"session","id":"session-appendable"}`, }, { name: "old-history-only", line: `{"type":"message","message":{"role":"user","content":"old prompt"}}`, }, { name: "assistant-only", line: `{"type":"message","message":{"role":"assistant","content":"quoted "}}`, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { require.NoError(t, os.WriteFile(sessionPath, []byte(tt.line+"\n"), 0o600)) driver := &PiRPCDriver{} driver.launchCfg = cfg driver.sessionID = "session-appendable" got, err := driver.ReconcileRuntimeInput(t.Context(), "q1", driver.sessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationUnknown, got) }) } // A reconcile racing the append must be unknown before the marker is // materialized, then become acked once the complete user record exists. require.NoError(t, os.WriteFile(sessionPath, []byte(`{"type":"session","id":"session-appendable"}`+"\n"), 0o600)) driver := &PiRPCDriver{launchCfg: cfg, sessionID: "session-appendable"} beforeAppend, err := driver.ReconcileRuntimeInput(t.Context(), "q1", driver.sessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationUnknown, beforeAppend) file, err := os.OpenFile(sessionPath, os.O_APPEND|os.O_WRONLY, 0o600) require.NoError(t, err) _, err = file.WriteString(`{"type":"message","message":{"role":"user","content":"q1\n\n"}}` + "\n") require.NoError(t, err) require.NoError(t, file.Close()) afterAppend, err := driver.ReconcileRuntimeInput(t.Context(), "q1", driver.sessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationAcked, afterAppend) } func TestPiRPCReconcileUsesOnlyInputSpecificRejectionFenceForAbsent(t *testing.T) { stateDir := t.TempDir() cfg := drivercontract.LaunchConfig{ RuntimeType: runtimeconfig.TypePi, RuntimeStateDir: stateDir, WorkDir: t.TempDir(), } sessionDir := filepath.Join(stateDir, ".pi-sessions") require.NoError(t, os.MkdirAll(sessionDir, 0o700)) sessionPath := filepath.Join(sessionDir, "2026_session-negative-fence.jsonl") require.NoError(t, os.WriteFile(sessionPath, []byte(`{"type":"session","id":"session-negative-fence"}`+"\n"), 0o600)) driver := &PiRPCDriver{launchCfg: cfg, sessionID: "session-negative-fence"} got, err := driver.ReconcileRuntimeInput(t.Context(), "q1", driver.sessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationUnknown, got) driver.rememberDefinitiveInputRejection("q1", driver.sessionID) got, err = driver.ReconcileRuntimeInput(t.Context(), "q1", driver.sessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationAbsent, got) // A replacement has no access to the process-local rejection fence and must // not inherit authoritative absent from the same readable snapshot. replacement := &PiRPCDriver{launchCfg: cfg, sessionID: driver.sessionID} got, err = replacement.ReconcileRuntimeInput(t.Context(), "q1", replacement.sessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationUnknown, got) } func TestPiRPCDefinitivePromptRejectionProvidesAbsentFence(t *testing.T) { rejected := false binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{ {Cmd: "prompt", Resp: &rejected}, }) d := &PiRPCDriver{} cfg := piTestConfig(t, binary) result, err := d.Spawn(t.Context(), cfg) require.NoError(t, err) t.Cleanup(result.Cleanup) err = d.DeliverRuntimeInput(t.Context(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "rejected-prompt"}, "must be rejected before mutation"), result.SessionID) require.Error(t, err) got, err := d.ReconcileRuntimeInput(t.Context(), "rejected-prompt", result.SessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationAbsent, got) } func TestPiRPCPreflightCancellationProvidesAbsentFence(t *testing.T) { binary, _ := writeFakePiRPCModeBinary(t, nil) d := &PiRPCDriver{} cfg := piTestConfig(t, binary) result, err := d.Spawn(t.Context(), cfg) require.NoError(t, err) t.Cleanup(result.Cleanup) ctx, cancel := context.WithCancel(t.Context()) cancel() err = d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "preflight-canceled"}, "must be queued because no RPC bytes were written"), result.SessionID) require.ErrorIs(t, err, errRuntimeInputMutationNotAttempted) got, err := d.ReconcileRuntimeInput(t.Context(), "preflight-canceled", result.SessionID) require.NoError(t, err) require.Equal(t, RuntimeInputReconciliationAbsent, got) } // --------------------------------------------------------------------------- // Iteration-3 contract tests (busy steer, final usage/stop/error, compaction, // generation-safe settlement). // --------------------------------------------------------------------------- // TestPiRPCBusySteerImmediateReturnDelayedAck verifies DeliverAckedBusyRuntimeInput // returns as soon as the steer RPC is accepted, retains correlation until the // (delayed) marked user message is acknowledged exactly once, and the caller's // ownership stays pending-only in between. func TestPiRPCBusySteerImmediateReturnDelayedAck(t *testing.T) { dir := t.TempDir() barrier := filepath.Join(dir, "release") marker := "" scenario := []piFakeScript{ piFakePromptScript( piFakeAgentStart(), piFakeTurnStart(), ), piFakeSteerScript( piFakeUserMessage("steer-delayed", "steer\n\n"+marker), piFakeUserMessageEnd("steer-delayed", "steer\n\n"+marker), piFakeAgentEnd(false), piFakeAgentSettled(), ).withWaitFor(barrier), } binary, _ := writeFakePiRPCModeBinary(t, scenario) d := &PiRPCDriver{} ctx := context.Background() cfg := piTestConfig(t, binary) result, err := d.Spawn(ctx, cfg) require.NoError(t, err) defer result.Cleanup() require.NoError(t, d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "start"), result.SessionID)) // Root run active. require.Eventually(t, func() bool { d.mu.Lock() defer d.mu.Unlock() return d.rootRunActive }, 5*time.Second, 10*time.Millisecond) // Steer RPC accepted immediately; the marker is delayed behind the barrier. steerDone := make(chan error, 1) go func() { steerDone <- d.DeliverAckedBusyRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "steer-delayed"}, "steer"), result.SessionID) }() select { case err := <-steerDone: require.NoError(t, err, "steer returns on RPC acceptance, not acknowledgement") case <-time.After(5 * time.Second): t.Fatal("DeliverAckedBusyRuntimeInput did not return on steer acceptance") } // Correlation retained: not yet acked. d.mu.Lock() correlated := d.steerInFlight["steer-delayed"] d.mu.Unlock() require.True(t, correlated, "correlation retained until the marker is acknowledged") // Release the barrier so the marked message is acknowledged exactly once. require.NoError(t, os.WriteFile(barrier, []byte("go"), 0o600)) events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle) acked := 0 for _, e := range events { if e.kind == consts.ParsedEventKindRuntimeInputAcked && e.runtimeInputID == "steer-delayed" { acked++ } } require.Equal(t, 1, acked, "delayed marker finalizes exactly once") // Correlation cleared after acknowledgement. d.mu.Lock() correlated = d.steerInFlight["steer-delayed"] d.mu.Unlock() require.False(t, correlated, "correlation cleared after acknowledgement") } // TestPiRPCDuplicateAgentEndAssistantSingleUsage verifies the normal // source-faithful path: an ordinary assistant message_end carrying usage, then // agent_end.messages repeating the SAME assistant, then agent_settled. Exactly // one Usage with the stable invocation id must be emitted. func TestPiRPCDuplicateAgentEndAssistantSingleUsage(t *testing.T) { final := piFakeUsage(30, 12, 0, 42) assistant := piFakeAssistantMessage( []map[string]any{piFakeTextBlock("hello")}, final, ) scenario := []piFakeScript{ piFakePromptScript( piFakeAgentStart(), piFakeTurnStart(), piFakeAssistantStart(nil, nil), piFakeTextDelta(0, "hello", final), piFakeAssistantEnd([]map[string]any{piFakeTextBlock("hello")}, final), piFakeTurnEnd([]map[string]any{piFakeTextBlock("hello")}, final), // agent_end.messages repeats the SAME assistant (normal pi behavior). piFakeAgentEnd(false, assistant), piFakeAgentSettled(), ), } binary, _ := writeFakePiRPCModeBinary(t, scenario) d := &PiRPCDriver{} ctx := context.Background() cfg := piTestConfig(t, binary) result, err := d.Spawn(ctx, cfg) require.NoError(t, err) defer result.Cleanup() require.NoError(t, d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "x"), result.SessionID)) events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle) require.Equal(t, 1, countKind(events, consts.ParsedEventKindUsage), "duplicate agent_end assistant must not double-emit Usage") require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd), "exactly one root TurnEnd") for _, e := range events { if e.kind == consts.ParsedEventKindUsage { require.Equal(t, uint64(30), e.inputTokens, "final usage from message_end") require.Equal(t, uint64(42), e.activeContextTokens) require.NotEmpty(t, e.invocationID, "stable invocation id") } } } // TestPiRPCTwoAssistantInvocationsTwoUsage verifies dedup is identity-based, // not overbroad: two distinct assistant invocations in one root generation each // emit one Usage. func TestPiRPCTwoAssistantInvocationsTwoUsage(t *testing.T) { usageA := piFakeUsage(10, 5, 0, 15) usageB := piFakeUsage(20, 9, 0, 29) assistantA := piFakeAssistantMessage([]map[string]any{piFakeTextBlock("first")}, usageA) assistantB := piFakeAssistantMessage([]map[string]any{piFakeTextBlock("second")}, usageB) assistantA["responseId"] = "resp-a" assistantB["responseId"] = "resp-b" scenario := []piFakeScript{ piFakePromptScript( piFakeAgentStart(), piFakeTurnStart(), piFakeAssistantStart(nil, nil), piFakeTextDelta(0, "first", usageA), piFakeAssistantEnd([]map[string]any{piFakeTextBlock("first")}, usageA), piFakeTurnEnd([]map[string]any{piFakeTextBlock("first")}, usageA), // A second genuine assistant invocation in the same root run. piFakeTurnStart(), piFakeAssistantStart(nil, nil), piFakeTextDelta(0, "second", usageB), piFakeAssistantEnd([]map[string]any{piFakeTextBlock("second")}, usageB), piFakeTurnEnd([]map[string]any{piFakeTextBlock("second")}, usageB), piFakeAgentEnd(false, assistantA, assistantB), piFakeAgentSettled(), ), } binary, _ := writeFakePiRPCModeBinary(t, scenario) d := &PiRPCDriver{} ctx := context.Background() cfg := piTestConfig(t, binary) result, err := d.Spawn(ctx, cfg) require.NoError(t, err) defer result.Cleanup() require.NoError(t, d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "x"), result.SessionID)) events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle) require.Equal(t, 2, countKind(events, consts.ParsedEventKindUsage), "two distinct assistant invocations -> two Usage") ids := map[string]bool{} for _, e := range events { if e.kind == consts.ParsedEventKindUsage { require.NotEmpty(t, e.invocationID) ids[e.invocationID] = true } } require.Len(t, ids, 2, "distinct stable invocation ids") } // TestPiRPCCompletionOnlyUsage verifies a provider that keeps message_update.usage // zero until completion still produces exactly one Usage from the final // assistant message.usage. func TestPiRPCCompletionOnlyUsage(t *testing.T) { zero := piFakeUsage(0, 0, 0, 0) final := piFakeUsage(42, 18, 7, 67) scenario := []piFakeScript{ piFakePromptScript( piFakeAgentStart(), piFakeTurnStart(), piFakeAssistantStart(nil, nil), piFakeTextDelta(0, "hello", zero), piFakeAssistantEnd([]map[string]any{piFakeTextBlock("hello")}, final), piFakeTurnEnd([]map[string]any{piFakeTextBlock("hello")}, final), piFakeAgentEnd(false), piFakeAgentSettled(), ), } binary, _ := writeFakePiRPCModeBinary(t, scenario) d := &PiRPCDriver{} ctx := context.Background() cfg := piTestConfig(t, binary) result, err := d.Spawn(ctx, cfg) require.NoError(t, err) defer result.Cleanup() require.NoError(t, d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "x"), result.SessionID)) events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle) require.Equal(t, 1, countKind(events, consts.ParsedEventKindUsage)) for _, e := range events { if e.kind == consts.ParsedEventKindUsage { require.Equal(t, uint64(42), e.inputTokens, "final authoritative usage input") require.Equal(t, uint64(18), e.outputTokens) require.Equal(t, uint64(7), e.cachedInputTokens) require.Equal(t, uint64(67), e.activeContextTokens) require.NotEmpty(t, e.invocationID) } } } // TestPiRPCRepeatedStreamingUsagePlusFinal verifies repeated message_update.usage // snapshots produce one Usage with the latest cumulative final value. func TestPiRPCRepeatedStreamingUsagePlusFinal(t *testing.T) { final := piFakeUsage(10, 5, 0, 15) scenario := []piFakeScript{ piFakePromptScript( piFakeAgentStart(), piFakeTurnStart(), piFakeAssistantStart(nil, nil), piFakeTextDelta(0, "a", piFakeUsage(1, 1, 0, 2)), piFakeTextDelta(0, "b", piFakeUsage(5, 2, 0, 7)), piFakeTextDelta(0, "c", piFakeUsage(9, 4, 0, 13)), piFakeAssistantEnd([]map[string]any{piFakeTextBlock("abc")}, final), piFakeTurnEnd([]map[string]any{piFakeTextBlock("abc")}, final), piFakeAgentEnd(false), piFakeAgentSettled(), ), } binary, _ := writeFakePiRPCModeBinary(t, scenario) d := &PiRPCDriver{} ctx := context.Background() cfg := piTestConfig(t, binary) result, err := d.Spawn(ctx, cfg) require.NoError(t, err) defer result.Cleanup() require.NoError(t, d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "x"), result.SessionID)) events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle) require.Equal(t, 1, countKind(events, consts.ParsedEventKindUsage), "one Usage per assistant invocation") for _, e := range events { if e.kind == consts.ParsedEventKindUsage { require.Equal(t, uint64(10), e.inputTokens, "final cumulative input") require.Equal(t, uint64(15), e.activeContextTokens) } } } // TestPiRPCMissingMessageEndDerivesFromAgentEnd verifies that when the normal // assistant message_end is absent (early provider failure), agent_end.messages // supplies the authoritative final usage/stop/error. func TestPiRPCMissingMessageEndDerivesFromAgentEnd(t *testing.T) { final := piFakeUsage(30, 12, 0, 42) scenario := []piFakeScript{ piFakePromptScript( piFakeAgentStart(), // No assistant message_end: the run fails early. piFakeAgentEnd(false, piFakeAssistantMessage( []map[string]any{piFakeTextBlock("partial")}, final, )), piFakeAgentSettled(), ), } binary, _ := writeFakePiRPCModeBinary(t, scenario) d := &PiRPCDriver{} ctx := context.Background() cfg := piTestConfig(t, binary) result, err := d.Spawn(ctx, cfg) require.NoError(t, err) defer result.Cleanup() require.NoError(t, d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "x"), result.SessionID)) events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle) require.Equal(t, 1, countKind(events, consts.ParsedEventKindUsage), "usage derived from agent_end.messages") for _, e := range events { if e.kind == consts.ParsedEventKindUsage { require.Equal(t, uint64(30), e.inputTokens) require.Equal(t, uint64(42), e.activeContextTokens) } } require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd)) } // TestPiRPCAbortBeforeAgentStart verifies abort before any agent_start produces // a settled root with stopReason "aborted" (not end_turn). func TestPiRPCAbortBeforeAgentStart(t *testing.T) { scenario := []piFakeScript{ piFakePromptScript(), // response only; no agent_start yet piFakeAbortScript( piFakeAgentEnd(false), piFakeAgentSettled(), ), } binary, _ := writeFakePiRPCModeBinary(t, scenario) d := &PiRPCDriver{} ctx := context.Background() cfg := piTestConfig(t, binary) result, err := d.Spawn(ctx, cfg) require.NoError(t, err) defer result.Cleanup() require.NoError(t, d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "x"), result.SessionID)) // Root active after prompt acceptance (no agent_start yet). d.mu.Lock() active := d.rootRunActive d.mu.Unlock() require.True(t, active) abortCtx, cancel := context.WithTimeout(ctx, 10*time.Second) defer cancel() require.NoError(t, d.InterruptActiveRun(abortCtx, result.SessionID)) events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle) require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd)) for _, e := range events { if e.kind == consts.ParsedEventKindTurnEnd { require.Equal(t, "aborted", e.stopReason, "abort-before-agent_start stop reason") } } } // TestPiRPCRetrySuccessClearsError verifies a transient error followed by a // successful retry produces NO final Error at settlement. func TestPiRPCRetrySuccessClearsError(t *testing.T) { usage := piFakeUsage(10, 5, 0, 15) scenario := []piFakeScript{ piFakePromptScript( piFakeAgentStart(), piFakeTurnStart(), piFakeAssistantStart(nil, nil), piFakeTextDelta(0, "first attempt", usage), piFakeAssistantEnd([]map[string]any{piFakeTextBlock("first attempt")}, usage), piFakeTurnEnd([]map[string]any{piFakeTextBlock("first attempt")}, usage), piFakeAgentEnd(true, piFakeAssistantMessage([]map[string]any{piFakeTextBlock("first attempt")}, usage)), piFakeAutoRetryStart(1, 3), piFakeAgentStart(), piFakeTurnStart(), piFakeAssistantStart(nil, nil), piFakeTextDelta(0, "retried ok", usage), piFakeAssistantEnd([]map[string]any{piFakeTextBlock("retried ok")}, usage), piFakeTurnEnd([]map[string]any{piFakeTextBlock("retried ok")}, usage), piFakeAutoRetryEnd(true, 1), piFakeAgentEnd(false), piFakeAgentSettled(), ), } binary, _ := writeFakePiRPCModeBinary(t, scenario) d := &PiRPCDriver{} ctx := context.Background() cfg := piTestConfig(t, binary) result, err := d.Spawn(ctx, cfg) require.NoError(t, err) defer result.Cleanup() require.NoError(t, d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "retry"), result.SessionID)) events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle) require.Equal(t, 0, countKind(events, consts.ParsedEventKindError), "successful retry clears prior error") require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd)) } // TestPiRPCFinalRetryFailure verifies a final retry failure updates the root // error outcome and emits exactly one nonretrying Error at settlement. func TestPiRPCFinalRetryFailure(t *testing.T) { usage := piFakeUsage(10, 5, 0, 15) scenario := []piFakeScript{ piFakePromptScript( piFakeAgentStart(), piFakeTurnStart(), piFakeAssistantStart(nil, nil), piFakeTextDelta(0, "boom", usage), piFakeAssistantEnd([]map[string]any{piFakeTextBlock("boom")}, usage), piFakeTurnEnd([]map[string]any{piFakeTextBlock("boom")}, usage), piFakeAgentEnd(true), piFakeAutoRetryStart(1, 2), piFakeAgentStart(), piFakeTurnStart(), piFakeAssistantStart(nil, nil), piFakeTextDelta(0, "boom again", usage), piFakeAssistantEnd([]map[string]any{piFakeTextBlock("boom again")}, usage), piFakeTurnEnd([]map[string]any{piFakeTextBlock("boom again")}, usage), piFakeAutoRetryEnd(false, 2), piFakeAgentEnd(false), piFakeAgentSettled(), ), } binary, _ := writeFakePiRPCModeBinary(t, scenario) d := &PiRPCDriver{} ctx := context.Background() cfg := piTestConfig(t, binary) result, err := d.Spawn(ctx, cfg) require.NoError(t, err) defer result.Cleanup() require.NoError(t, d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "retry-fail"), result.SessionID)) events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle) require.Equal(t, 1, countKind(events, consts.ParsedEventKindError), "final retry failure emits one Error") require.Equal(t, 1, countKind(events, consts.ParsedEventKindTurnEnd)) } // TestPiRPCCompactionFrames covers successful, aborted, retry, and final-failed // compaction_end semantics. func TestPiRPCCompactionFrames(t *testing.T) { tests := []struct { name string scenario []piFakeScript expectErr int expectConfirmed int expectTurnEnd int }{ { name: "success", scenario: []piFakeScript{piFakePromptScript( piFakeAgentStart(), piFakeTurnStart(), piFakeCompactionStart("manual"), piFakeCompactionEnd("manual"), piFakeAgentEnd(false), piFakeAgentSettled(), )}, expectErr: 0, expectConfirmed: 1, expectTurnEnd: 1, }, { name: "aborted", scenario: []piFakeScript{piFakePromptScript( piFakeAgentStart(), piFakeTurnStart(), piFakeCompactionStart("manual"), piFakeCompactionEndAborted(), piFakeAgentEnd(false), piFakeAgentSettled(), )}, expectErr: 0, expectConfirmed: 0, expectTurnEnd: 1, }, { name: "failed-with-retry", scenario: []piFakeScript{piFakePromptScript( piFakeAgentStart(), piFakeTurnStart(), piFakeCompactionStart("overflow"), piFakeCompactionEndRetry(), piFakeAgentEnd(false), piFakeAgentSettled(), )}, expectErr: 0, expectConfirmed: 0, expectTurnEnd: 1, }, { name: "final-failed", scenario: []piFakeScript{piFakePromptScript( piFakeAgentStart(), piFakeTurnStart(), piFakeCompactionStart("threshold"), piFakeCompactionEndFailed(), piFakeAgentEnd(false), piFakeAgentSettled(), )}, expectErr: 1, expectConfirmed: 0, expectTurnEnd: 1, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { binary, _ := writeFakePiRPCModeBinary(t, tt.scenario) d := &PiRPCDriver{} ctx := context.Background() cfg := piTestConfig(t, binary) result, err := d.Spawn(ctx, cfg) require.NoError(t, err) defer result.Cleanup() require.NoError(t, d.DeliverRuntimeInput(ctx, testRuntimeInputWithConfig(drivercontract.RuntimeInput{}, "compact"), result.SessionID)) events := collectPiEvents(t, result, consts.ParsedEventKindRuntimeIdle) require.Equal(t, tt.expectErr, countKind(events, consts.ParsedEventKindError), "error count") require.Equal(t, tt.expectTurnEnd, countKind(events, consts.ParsedEventKindTurnEnd), "TurnEnd count") // confirmed compaction metadata present exactly when expected confirmed := 0 for _, e := range events { if e.kind == consts.ParsedEventKindThinking && e.text == "Pi context compaction completed" { confirmed++ } } require.Equal(t, tt.expectConfirmed, confirmed, "confirmed compaction text") }) } }