package agent import ( "context" "encoding/json" "errors" "fmt" "io" "log/slog" "os" "os/exec" "path/filepath" "strings" "sync" "time" "github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract" "github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig" "github.com/tidbcloud/AgentTeam/desktop/internal/consts" "github.com/tidbcloud/AgentTeam/internal/core/rewindcaps" ) // PiRPCRecoveryStrategy is the rewind strategy used by the Pi RPC driver. It // restores a Pi session by truncating its persisted JSONL tree at the entry // that was the current leaf before the AgentTeam message being rewound. const PiRPCRecoveryStrategy = rewindcaps.StrategyPiSessionTreeRewind const ( // piRPCMaxLineBytes bounds a single JSONL frame from the pi process. piRPCMaxLineBytes = 32 * 1024 * 1024 // piRPCActiveTurnPollTimeout bounds how long the driver waits for the // authoritative agent_settled boundary after an explicit abort before // classifying the outcome. piRPCActiveTurnPollTimeout = 60 * time.Second // piRPCRequestTimeout bounds a single RPC request round-trip. piRPCRequestTimeout = 300 * time.Second ) // PiRPCDriver drives the pi CLI in headless RPC mode (`pi --mode rpc`). // // Pi RPC is a long-lived subprocess speaking strict LF-delimited JSONL over // stdin/stdout. The driver spawns one process per AgentTeam Process, delivers // messages as `prompt`/`steer` commands, and translates the streamed event // protocol into drivercontract.ParsedEvent values. // // Root lifecycle (source-faithful to pi-mono commit // eb1f87fa9a29e27e0c63dcb40dbed9a3624c82b1): // - A Pi `turn_start`/`turn_end` is one assistant/tool cycle. // - A Pi `agent_start`/`agent_end` is one low-level run that can be followed // by an automatic retry, an overflow compaction, or queued steer/follow-up. // - `agent_settled` is the ONLY authoritative fully-settled boundary. It // produces exactly one AgentTeam root ParsedEventKindTurnEnd (with the final // assistant stop/error/aborted reason) and the appropriate RuntimeIdle. // // The driver keeps a generation-fenced root run: a successful `prompt` RPC // response returns command acceptance to the caller while the root run stays // active (rootRunActive=true) until agent_settled. Abort and interrupting // replacement wait for the same generation's agent_settled. type PiRPCDriver struct { mu sync.Mutex writeMu sync.Mutex stdin io.WriteCloser lines chan string output *piTranslatedOutput nextID int pending map[int]chan piRPCFrame sessionID string launchCfg drivercontract.LaunchConfig sessionEmitted bool runtimeCtx context.Context transportClosed bool // --- command acceptance (per DeliverRuntimeInput) --- // promptAccepted resolves on the prompt RPC preflight response. promptAccepted chan error promptInputID string // --- root run (the authoritative AgentTeam root turn) --- rootRunGeneration uint64 rootRunTurnID string // stable turn identity for ActiveTurnIDProvider / Process fence rootRunActive bool rootRunSettled chan struct{} // closed once agent_settled settles this generation rootStartEmitted bool // authoritative root-start metadata once per generation rootStopReason string // final assistant stop reason captured at settlement rootAborted bool // an abort was dispatched for this generation rootError string // final provider/retry error acked exactly once rootErrorEmitted bool // the root error was already projected as an Error event // --- low-level agent run (agent_start .. agent_end) --- agentRunActive bool agentRunCount int // --- turn cycle (turn_start .. turn_end) --- turnActive bool turnCount int contextCompacted bool // Native Pi compaction can finish outside a root agent run. Keep its // ownership separate from rootRunActive so a late compaction_end from an // already-settled root cannot manufacture a second RuntimeIdle boundary. nativeCompactionInProgress bool nativeCompactionRootGeneration uint64 // --- steer correlation (registered before send, cleared on ack/retire) --- steerInFlight map[string]bool // startupSteerIdleObservedAt delays negative reconciliation for a steer that // Pi accepted before an agent run became active. Pi's steer response and its // in-memory user entry are not atomic: get_state/get_entries can briefly // observe idle+absent immediately before the marked entry is committed. startupSteerIdleObservedAt map[string]time.Time // definitiveInputRejections records input-specific negative fences from // the current driver generation. A rejected RPC response (or a preflight // failure before any bytes were written) is the only Pi-side evidence that // permits reconciliation to return authoritative absent. This state is // intentionally process-local: a replacement driver has no such fence and // must treat a readable transcript without the marker as unknown. definitiveInputRejections map[string]string // --- assistant message streaming dedup state (reset per root run) --- assistant *piAssistantStream // ackedInputIDs dedupes runtime_input_acked emissions per generation so // message_start + message_end both carrying the marker produce exactly one. ackedInputIDs map[string]bool // reconciledAckedInputIDs retains positive evidence read from Pi's live // get_entries snapshot without suppressing the independent streamed marker. reconciledAckedInputIDs map[string]bool // initialAckWaiters are used by Spawn's first correlated prompt. The // prompt RPC response proves command acceptance, not that the user message // crossed Pi's durable session/model input boundary. initialAckWaiters map[string]chan error // finalizedAssistantIDs tracks the stable ModelInvocationID (responseId or // session:assistantTimestamp) of assistant invocations already finalized in // this root generation. Duplicate assistant snapshots (message_end + // agent_end.messages carrying the same assistant) must never re-emit Usage / // text / tool projections for an already-finalized identity. Cleared only // when a new root generation starts. finalizedAssistantIDs map[string]bool entryIDAtLastInput string } var errPiRPCTransportClosed = fmt.Errorf("pi RPC transport closed") // errPiRPCTransportAmbiguous classifies a steer/prompt whose transport outcome // is unknown after a write: the runtime may have accepted the frame before the // pipe disconnected. Process treats this as fail-closed correlated ownership // (no retry) rather than a definite rejection. var errPiRPCTransportAmbiguous = fmt.Errorf("pi RPC transport outcome is ambiguous") // piRPCMutationAmbiguous reports whether a mutating RPC outcome is ambiguous. func piRPCMutationAmbiguous(err error) bool { if errors.Is(err, errRuntimeInputMutationNotAttempted) { return false } return errors.Is(err, errPiRPCTransportAmbiguous) || errors.Is(err, errPiRPCTransportClosed) || errors.Is(err, io.ErrClosedPipe) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) } func (d *PiRPCDriver) Name() string { return "pi" } // SupportsStdinNotification reports that pi accepts RPC `prompt` commands on // stdin. func (d *PiRPCDriver) SupportsStdinNotification() bool { return true } // MCPToolPrefix returns the runtime-specific MCP tool prefix. Pi's MCP bridge // registers tools with a sanitized name, so AgentTeam exposes agentteam tools using the // same namespace prefix as other JSON-line runtimes. func (d *PiRPCDriver) MCPToolPrefix() string { return "mcp__agentteam__" } // UsesExplicitTurnLifecycle reports that pi's authoritative liveness boundary // is agent_settled; no Pi turn/agent event independently makes the root turn // terminal. func (d *PiRPCDriver) UsesExplicitTurnLifecycle() bool { return true } // SupportsAckedBusyRuntimeInputDelivery reports that pi can queue a correlated // steer message while a root run is active and later acknowledge it with the same // id. func (d *PiRPCDriver) SupportsAckedBusyRuntimeInputDelivery() bool { return true } // SupportsCorrelatedRuntimeInputDelivery reports that pi echoes the runtime // input id marker when a message enters its durable session transcript. func (d *PiRPCDriver) SupportsCorrelatedRuntimeInputDelivery() bool { return true } // SupportsBusyStdinDelivery reports that pi can accept a routed message while // a root run is active via the steer queue. func (d *PiRPCDriver) SupportsBusyStdinDelivery() bool { return true } // --------------------------------------------------------------------------- // Spawn // --------------------------------------------------------------------------- // Spawn launches pi in RPC mode and delivers the initial prompt. The pinned // Pi protocol has no deterministic pre-input invalid-session authority: its // get_state preflight is a read-only consistency check, so every failure is // fail-closed and no fresh-session retry is attempted. func (d *PiRPCDriver) Spawn(ctx context.Context, cfg drivercontract.LaunchConfig) (drivercontract.SpawnResult, error) { return d.spawnOnce(ctx, cfg) } func (d *PiRPCDriver) spawnOnce(ctx context.Context, cfg drivercontract.LaunchConfig) (drivercontract.SpawnResult, error) { carrier, err := cfg.RequireCompiledPrompt() if err != nil { return drivercontract.SpawnResult{}, piStartupFailure(err) } cfg.CompiledPrompt = carrier if err := materializeContractCarrier(cfg); err != nil { return drivercontract.SpawnResult{}, piStartupFailure(fmt.Errorf("pi: materialize contract carrier: %w", err)) } sessionID, resumed, err := resolvePiLaunchSession(cfg) if err != nil { return drivercontract.SpawnResult{}, piStartupFailure(err) } cfg.SessionID = sessionID prepared, err := preparePiLaunch(cfg, sessionID) if err != nil { return drivercontract.SpawnResult{}, piStartupFailure(err) } translatedReader, translatedWriter := io.Pipe() output := newPiTranslatedOutput(512) output.attachWriter(translatedWriter) lines := output.lines readerDone := make(chan struct{}) go func() { defer close(readerDone) defer translatedWriter.Close() for line := range lines { if _, err := io.WriteString(translatedWriter, line+"\n"); err != nil { return } } }() cmd, stdin, stdout, stderr, terminate, wait, err := startPiRPCProcess(ctx, cfg, prepared) if err != nil { prepared.Cleanup() output.close() _ = translatedReader.Close() return drivercontract.SpawnResult{}, piStartupFailureWithStderr(err, stderr) } d.initializeDriverState(cfg, sessionID, stdin, lines, output, ctx) cleanupOnce := sync.Once{} cleanup := func() { cleanupOnce.Do(func() { // closeTransport emits the terminal Error/TurnEnd/RuntimeIdle for a // transport failure, so it must run while the output queue is still // open. Then abort the queue and tear down the process. d.closeTransport() output.close() terminate() }) } preparedCleanupOnce := sync.Once{} cleanupPrepared := func() { preparedCleanupOnce.Do(prepared.Cleanup) } cleanupAll := func() { cleanup() cleanupPrepared() } cleanupFailedSpawn := func() { cleanup() _ = translatedReader.Close() _ = wait() <-readerDone cleanupPrepared() } go d.readLoop(stdout, cleanup) d.emit(translatedEvent{Kind: consts.ParsedEventKindSessionInit, SessionID: sessionID}) d.mu.Lock() d.sessionEmitted = true d.mu.Unlock() if resumed { // get_state is a read-only consistency check. The pinned Pi response // contract has no typed pre-input invalid-session authority, so every // failure must stop this process before any prompt bytes are written. if err := d.preflightResumedSession(ctx, sessionID); err != nil { cleanupFailedSpawn() return drivercontract.SpawnResult{}, fmt.Errorf("pi: session preflight: %w", err) } } preparedRecoveryPoint := prepareSpawnRecoveryPoint(d, cfg.InitialDeliveryMessageID, sessionID) initialPrompt := piInitialPrompt(cfg, resumed) if strings.TrimSpace(initialPrompt) != "" || len(cfg.PromptAttachments) > 0 { // Deliver the initial prompt synchronously so the preflight is accepted // and the generation-fenced root run is active before Spawn returns. The // delivery goroutine inside DeliverRuntimeInput would otherwise race the // first acked-busy steer (or any other RPC) for the stdin write; when the // steer wins, it is written before the prompt and the runtime never // receives the correlated turn the driver is waiting to settle. input := drivercontract.RuntimeInput{ ID: cfg.InitialDeliveryMessageID, Resume: resumed, CompiledPrompt: carrier, Attachments: cfg.PromptAttachments, } if err := d.DeliverRuntimeInput(ctx, input, sessionID); err != nil { initialPromptErr := fmt.Errorf("pi: initial RPC prompt: %w", err) cleanupFailedSpawn() return drivercontract.SpawnResult{}, piStartupFailureWithStderr(initialPromptErr, stderr) } if inputID := strings.TrimSpace(input.ID); inputID != "" { if err := d.waitForInitialAck(ctx, inputID); err != nil { initialPromptErr := fmt.Errorf("pi: initial RPC prompt acknowledgement: %w", err) cleanupFailedSpawn() return drivercontract.SpawnResult{}, piStartupFailureWithStderr(initialPromptErr, stderr) } } } return drivercontract.SpawnResult{ Cmd: cmd, Stdin: stdin, Stdout: translatedReader, StderrTail: stderr.Tail, SessionID: sessionID, SessionResumed: resumed, Terminate: cleanup, Wait: func() error { err := wait() cleanupAll() <-readerDone return err }, Cleanup: cleanupAll, PreparedRecoveryPoint: preparedRecoveryPoint, }, nil } // initializeDriverState resets per-launch driver state under the driver lock. func (d *PiRPCDriver) initializeDriverState( cfg drivercontract.LaunchConfig, sessionID string, stdin io.WriteCloser, lines chan string, output *piTranslatedOutput, ctx context.Context, ) { d.mu.Lock() d.stdin = stdin d.lines = lines d.output = output d.pending = make(map[int]chan piRPCFrame) d.steerInFlight = make(map[string]bool) d.startupSteerIdleObservedAt = make(map[string]time.Time) d.definitiveInputRejections = make(map[string]string) d.sessionID = sessionID d.launchCfg = clonePiLaunchConfig(cfg) d.sessionEmitted = false d.rootRunGeneration = 0 d.rootRunTurnID = "" d.rootRunActive = false d.rootRunSettled = closedPiSignal() d.rootStartEmitted = false d.rootStopReason = "" d.rootAborted = false d.rootError = "" d.rootErrorEmitted = false d.agentRunActive = false d.agentRunCount = 0 d.turnActive = false d.turnCount = 0 d.contextCompacted = false d.nativeCompactionInProgress = false d.nativeCompactionRootGeneration = 0 d.promptAccepted = nil d.promptInputID = "" d.transportClosed = false d.runtimeCtx = ctx d.assistant = nil d.ackedInputIDs = make(map[string]bool) d.reconciledAckedInputIDs = make(map[string]bool) d.initialAckWaiters = make(map[string]chan error) d.finalizedAssistantIDs = make(map[string]bool) d.entryIDAtLastInput = "" d.mu.Unlock() } // startPiRPCProcess builds and starts the pi child process. func startPiRPCProcess( ctx context.Context, cfg drivercontract.LaunchConfig, prepared piPreparedLaunch, ) (*exec.Cmd, io.WriteCloser, io.ReadCloser, *capturingPrefixWriter, func(), func() error, error) { binary := strings.TrimSpace(cfg.Binary) if binary == "" { binary = "pi" } command, args := runtimeconfig.RuntimeCommand(runtimeconfig.TypePi, binary, prepared.Args) cmd := exec.CommandContext(ctx, command, args...) cmd.Dir = prepared.WorkDir cmd.Env = prepared.Env stderr := newCapturingPrefixWriter(fmt.Sprintf("[pi %s] ", cfg.AgentID), 32*1024) cmd.Stderr = stderr configureRuntimeCommand(cmd) stdout, err := cmd.StdoutPipe() if err != nil { return nil, nil, nil, stderr, nil, nil, fmt.Errorf("pi: get stdout pipe: %w", err) } stdin, err := cmd.StdinPipe() if err != nil { _ = stdout.Close() return nil, nil, nil, stderr, nil, nil, fmt.Errorf("pi: get stdin pipe: %w", err) } if err := cmd.Start(); err != nil { _ = stdout.Close() _ = stdin.Close() return nil, nil, nil, stderr, nil, nil, fmt.Errorf("pi: start RPC process: %w", err) } terminateOnce := sync.Once{} terminate := func() { terminateOnce.Do(func() { _ = stdin.Close() _ = stdout.Close() terminateRuntimeCommand(cmd) }) } return cmd, stdin, stdout, stderr, terminate, cmd.Wait, nil } // --------------------------------------------------------------------------- // Recovery // --------------------------------------------------------------------------- // PrepareRecoveryPoint captures the current session leaf as a rewind boundary. func (d *PiRPCDriver) PrepareRecoveryPoint() (string, string, bool, error) { d.mu.Lock() sessionID := d.sessionID cfg := clonePiLaunchConfig(d.launchCfg) d.mu.Unlock() if sessionID == "" { return "", "", false, nil } boundary, err := d.currentLeafEntryID(cfg, sessionID) if err != nil { return "", "", false, err } boundaryKind := rewindcaps.RecoveryBoundaryResumeAt if strings.TrimSpace(boundary) == "" { boundaryKind = rewindcaps.RecoveryBoundaryBeforeFirstInput } ref, err := rewindcaps.EncodePiSessionTreeRewindRef(sessionID, boundary, boundaryKind) if err != nil { return "", "", false, err } return PiRPCRecoveryStrategy, ref, true, nil } // PrepareRecoveryPointForMessage captures the session leaf that was current // before the AgentTeam message identified by messageID entered the transcript. func (d *PiRPCDriver) PrepareRecoveryPointForMessage(messageID string) (string, string, bool, error) { messageID = strings.TrimSpace(messageID) if messageID == "" { return "", "", false, nil } d.mu.Lock() sessionID := d.sessionID cfg := clonePiLaunchConfig(d.launchCfg) d.mu.Unlock() if sessionID == "" { return "", "", false, nil } entryID := "" if path, err := piSessionFilePath(cfg, sessionID); err == nil { if leaf, leafErr := piSessionLeafID(path); leafErr == nil && leaf != "" { entryID = leaf } } boundaryKind := rewindcaps.RecoveryBoundaryResumeAt if strings.TrimSpace(entryID) == "" { boundaryKind = rewindcaps.RecoveryBoundaryBeforeFirstInput } ref, err := rewindcaps.EncodePiSessionTreeRewindRef(sessionID, entryID, boundaryKind) if err != nil { return "", "", false, err } return PiRPCRecoveryStrategy, ref, true, nil } // RestoreRecoveryPoint is unsupported for live restore. Pi rewind restarts a // process from the truncated durable session file (detached restore). func (d *PiRPCDriver) RestoreRecoveryPoint(_ context.Context, strategy string, recoveryRef string) error { trimmedStrategy := strings.TrimSpace(strategy) if trimmedStrategy != PiRPCRecoveryStrategy { return fmt.Errorf("pi: unsupported recovery strategy %q", trimmedStrategy) } if strings.TrimSpace(recoveryRef) == "" { return fmt.Errorf("pi: live recovery restore requires a recovery ref") } return fmt.Errorf("pi: live recovery restore is not supported; use detached restore") } // RestoreRecoveryPointDetached truncates the Pi session JSONL tree at the // boundary entry so a later launch resumes without later context. func (d *PiRPCDriver) RestoreRecoveryPointDetached( ctx context.Context, cfg drivercontract.LaunchConfig, strategy string, recoveryRef string, ) error { if err := ctx.Err(); err != nil { return err } if strings.TrimSpace(strategy) != PiRPCRecoveryStrategy { return fmt.Errorf("pi: unsupported recovery strategy %q", strategy) } ref, err := rewindcaps.DecodePiSessionTreeRewindRef(recoveryRef) if err != nil { return fmt.Errorf("pi: decode recovery ref: %w", err) } if ref.SessionID == "" { return fmt.Errorf("pi: recovery ref is missing session id") } sessionPath, err := piSessionFilePath(cfg, ref.SessionID) if err != nil { return err } return piTruncateSessionAtBoundary(ctx, sessionPath, ref.BoundaryEntryID, ref.BoundaryKind) } // --------------------------------------------------------------------------- // RPC transport // --------------------------------------------------------------------------- func (d *PiRPCDriver) writeFrame(data []byte) error { d.writeMu.Lock() defer d.writeMu.Unlock() d.mu.Lock() stdin := d.stdin closed := d.transportClosed d.mu.Unlock() if stdin == nil || closed { return io.ErrClosedPipe } _, err := stdin.Write(data) return err } func (d *PiRPCDriver) forgetPending(id int) { d.mu.Lock() delete(d.pending, id) d.mu.Unlock() } func (d *PiRPCDriver) readLoop(stdout io.Reader, cleanup func()) { defer func() { cleanup() d.failPending(fmt.Errorf("pi RPC transport closed")) }() reader := newPiLineReader(stdout, piRPCMaxLineBytes) for { line, err := reader.Read() if err != nil { return } line = strings.TrimSpace(line) if line == "" { continue } d.handleLine(line) } } func (d *PiRPCDriver) handleLine(line string) { var frame piRPCFrame if err := json.Unmarshal([]byte(line), &frame); err != nil { slog.Warn("pi: invalid RPC frame", "error", err) return } if frame.Type == "response" { if len(frame.ID) > 0 { var id int if json.Unmarshal(frame.ID, &id) == nil { // A successful prompt response is the authoritative command // acceptance: start the generation-fenced root run synchronously // BEFORE routing the response, so the event frames that follow on // stdout observe rootRunActive=true and are not dropped as stale. if strings.TrimSpace(frame.Command) == "prompt" && frame.Success != nil && *frame.Success { d.beginRootRun() } d.mu.Lock() ch := d.pending[id] delete(d.pending, id) d.mu.Unlock() if ch != nil { ch <- frame } } } return } d.handleEvent(frame) } func (d *PiRPCDriver) failPending(err error) { d.mu.Lock() pending := d.pending d.pending = make(map[int]chan piRPCFrame) steerInFlight := d.steerInFlight d.steerInFlight = make(map[string]bool) d.startupSteerIdleObservedAt = make(map[string]time.Time) promptAccepted := d.promptAccepted d.promptAccepted = nil initialAckWaiters := d.initialAckWaiters d.initialAckWaiters = make(map[string]chan error) settled := d.rootRunSettled active := d.rootRunActive gen := d.rootRunGeneration turnID := d.rootRunTurnID stopReason := d.rootStopReason rootError := d.rootError aborted := d.rootAborted sessionID := d.sessionID if active { d.rootRunActive = false } d.promptInputID = "" d.mu.Unlock() for _, ch := range pending { success := false ch <- piRPCFrame{Type: "response", Success: &success, Error: err.Error()} } _ = steerInFlight if promptAccepted != nil { promptAccepted <- errPiRPCTransportClosed } for _, waiter := range initialAckWaiters { waiter <- err } // Transport retirement settles the active root run: enqueue the terminal // events (tagged with the completed generation) BEFORE releasing the // generation's waiter so a replacement can never consume an old idle after a // new handoff. if active { d.emitTerminal(gen, turnID, sessionID, stopReason, firstNonEmptyString(rootError, "pi RPC transport closed"), stopReason, aborted) d.settleRootRunLocked(settled) } } // emitTerminal enqueues the authoritative root terminal boundary for a // generation: a nonretrying Error (when the final outcome is a failure and it // was not already emitted), exactly one TurnEnd tagged with the completed turn // id, then RuntimeIdle. func (d *PiRPCDriver) emitTerminal(gen uint64, turnID string, sessionID string, stopReason string, errorMessage string, priorStopReason string, aborted bool) { completedTurnID := turnID if completedTurnID == "" { completedTurnID = fmt.Sprintf("pi-gen-%d", gen) } terminalMeta := map[string]any{runtimeCompletedTurnIDKey: completedTurnID} // A final failure emits exactly one nonretrying Error. If it was already // projected (auto_retry_end/compaction_end), do not emit a duplicate. if strings.TrimSpace(errorMessage) != "" { d.mu.Lock() alreadyEmitted := d.rootErrorEmitted d.rootErrorEmitted = true d.mu.Unlock() if !alreadyEmitted { d.emit(translatedEvent{Kind: consts.ParsedEventKindError, Message: errorMessage, SessionID: sessionID, Metadata: map[string]any{"will_retry": false}}) } } resolved := firstNonEmptyString(stopReason, priorStopReason) if aborted && resolved != "error" && resolved != consts.ParsedEventStopReasonEndTurn { // An abort was dispatched and no real assistant stop reason superseded // it. resolved = "aborted" } d.emit(translatedEvent{ Kind: consts.ParsedEventKindTurnEnd, SessionID: sessionID, StopReason: piStopReasonForLoop(resolved), Metadata: terminalMeta, }) d.emit(translatedEvent{Kind: consts.ParsedEventKindRuntimeIdle, SessionID: sessionID, Metadata: terminalMeta}) } func (d *PiRPCDriver) closeTransport() { d.mu.Lock() if d.transportClosed { d.mu.Unlock() return } d.transportClosed = true stdin := d.stdin d.stdin = nil promptAccepted := d.promptAccepted d.promptAccepted = nil initialAckWaiters := d.initialAckWaiters d.initialAckWaiters = make(map[string]chan error) steerInFlight := d.steerInFlight d.steerInFlight = make(map[string]bool) d.startupSteerIdleObservedAt = make(map[string]time.Time) settled := d.rootRunSettled active := d.rootRunActive gen := d.rootRunGeneration turnID := d.rootRunTurnID stopReason := d.rootStopReason rootError := d.rootError aborted := d.rootAborted sessionID := d.sessionID if active { d.rootRunActive = false } d.promptInputID = "" d.mu.Unlock() if stdin != nil { _ = stdin.Close() } if promptAccepted != nil { promptAccepted <- errPiRPCTransportClosed } for _, waiter := range initialAckWaiters { waiter <- errPiRPCTransportClosed } _ = steerInFlight // Transport failure settles the active root run: enqueue terminal events // (tagged with the completed generation) BEFORE releasing the waiter. if active { d.emitTerminal(gen, turnID, sessionID, stopReason, firstNonEmptyString(rootError, "pi RPC transport closed"), stopReason, aborted) d.settleRootRunLocked(settled) } } // settleRootRunLocked closes the generation's settled channel (idempotent). func (d *PiRPCDriver) settleRootRunLocked(settled chan struct{}) { if settled == nil { return } select { case <-settled: default: close(settled) } } func closedPiSignal() chan struct{} { ch := make(chan struct{}) close(ch) return ch } // --------------------------------------------------------------------------- // Output // --------------------------------------------------------------------------- func (d *PiRPCDriver) emit(event translatedEvent) { d.mu.Lock() output := d.output d.mu.Unlock() if output == nil { return } data, err := json.Marshal(event) if err != nil { return } output.emit(string(data)) } // ParseLine decodes a translated event line into a ParsedEvent. func (d *PiRPCDriver) ParseLine(line string) []drivercontract.ParsedEvent { var event translatedEvent if json.Unmarshal([]byte(strings.TrimSpace(line)), &event) != nil { return nil } parsed := drivercontract.ParsedEvent{ Kind: event.Kind, Text: event.Text, SectionID: event.SectionID, ToolName: event.ToolName, ToolCallID: event.ToolCallID, ToolInput: event.ToolInput, RuntimeInputID: event.RuntimeInputID, SessionID: event.SessionID, Message: event.Message, Metadata: event.Metadata, StopReason: event.StopReason, InputTokens: event.InputTokens, CachedInputTokens: event.CachedInputTokens, OutputTokens: event.OutputTokens, ActiveContextTokens: event.ActiveContextTokens, ContextWindowTokens: event.ContextWindowTokens, Model: event.Model, ModelInvocationID: event.ModelInvocationID, ContextExposureKnown: event.ContextExposureKnown, ContextExposureExcluded: event.ContextExposureExcluded, ContextCompacted: event.ContextCompacted, } switch event.Kind { case consts.ParsedEventKindSessionInit: if parsed.SessionID == "" { return nil } case consts.ParsedEventKindText, consts.ParsedEventKindThinking: if parsed.Text == "" && !parsedBoolMetadata(parsed.Metadata, codexAuthoritativeTurnStartMetadataKey) { return nil } case consts.ParsedEventKindRuntimeInputAcked: if parsed.RuntimeInputID == "" { return nil } case consts.ParsedEventKindToolCall, consts.ParsedEventKindToolOutput, consts.ParsedEventKindTurnEnd, consts.ParsedEventKindError, consts.ParsedEventKindUsage, consts.ParsedEventKindRuntimeCompacting, consts.ParsedEventKindSubagentActivity, consts.ParsedEventKindRuntimeIdle: default: return nil } return []drivercontract.ParsedEvent{parsed} } var ( _ drivercontract.Driver = (*PiRPCDriver)(nil) _ drivercontract.RuntimeInputDeliveryDriver = (*PiRPCDriver)(nil) _ drivercontract.AckedBusyRuntimeInputDriver = (*PiRPCDriver)(nil) _ drivercontract.InterruptingRuntimeInputDriver = (*PiRPCDriver)(nil) _ drivercontract.ActiveRunInterruptDriver = (*PiRPCDriver)(nil) _ drivercontract.ExplicitTurnLifecycleDriver = (*PiRPCDriver)(nil) _ drivercontract.CorrelatedRuntimeInputDeliveryDriver = (*PiRPCDriver)(nil) _ drivercontract.RecoveryDriver = (*PiRPCDriver)(nil) _ drivercontract.MessageBoundRecoveryDriver = (*PiRPCDriver)(nil) _ drivercontract.DetachedRecoveryDriver = (*PiRPCDriver)(nil) _ drivercontract.BusyStdinDeliveryDriver = (*PiRPCDriver)(nil) _ drivercontract.BusyStdinMessageDriver = (*PiRPCDriver)(nil) _ drivercontract.ActiveTurnIDProvider = (*PiRPCDriver)(nil) _ RuntimeInputReconciler = (*PiRPCDriver)(nil) ) func piRPCUserInput(input drivercontract.RuntimeInput) (any, error) { text := input.PromptText() if inputID := strings.TrimSpace(input.ID); inputID != "" { if text != "" { text += "\n\n" } text += piRuntimeInputMarker(inputID) } if len(input.Attachments) == 0 { return text, nil } parts := make([]map[string]any, 0, len(input.Attachments)+1) if strings.TrimSpace(text) != "" { parts = append(parts, map[string]any{"type": "text", "text": text}) } for _, attachment := range input.Attachments { uri := strings.TrimSpace(attachment.URI) if uri == "" { uri = strings.TrimSpace(attachment.Path) } if uri == "" { continue } mimeType := strings.ToLower(strings.TrimSpace(attachment.MIME)) switch { case strings.HasPrefix(mimeType, "image/"): parts = append(parts, map[string]any{"type": "image", "data": uri, "mimeType": mimeType}) default: parts = append(parts, map[string]any{"type": "text", "text": "Attachment: " + uri}) } } return parts, nil } func clonePiLaunchConfig(cfg drivercontract.LaunchConfig) drivercontract.LaunchConfig { cfg.Env = runtimeconfig.CloneStringMap(cfg.Env) return cfg } // piRuntimeInputMarkerPattern is the same durable marker convention AgentTeam uses // for Kimi; pi sessions persist user messages verbatim so the marker survives // restarts and can be matched against get_entries transcripts. var piRuntimeInputMarkerPattern = kimiRuntimeInputMarkerPattern func piRuntimeInputMarker(inputID string) string { return "" } func piRuntimeInputIDFromText(text string) string { match := piRuntimeInputMarkerPattern.FindStringSubmatch(strings.TrimSpace(text)) if len(match) != 2 { return "" } return strings.TrimSpace(match[1]) } // piRuntimeInputIDFromContent accepts both forms Pi uses for user content: // plain string content (the normal prompt/steer path) and an array of text or // image blocks (the attachment path). Keeping this extraction at the wire // boundary prevents the durable acknowledgement from depending on // whether a delivery happened to include an attachment. func piRuntimeInputIDFromContent(content json.RawMessage) string { if len(content) == 0 || strings.TrimSpace(string(content)) == "null" { return "" } var text string if json.Unmarshal(content, &text) == nil { return piRuntimeInputIDFromText(text) } for _, block := range decodePiContentBlocks(content) { if block.Type == "text" { if inputID := piRuntimeInputIDFromText(block.Text); inputID != "" { return inputID } } } return "" } // piSessionFilePath locates the durable session file for a session id. func piSessionFilePath(cfg drivercontract.LaunchConfig, sessionID string) (string, error) { sessionID = strings.TrimSpace(sessionID) if sessionID == "" { return "", fmt.Errorf("pi: session id is required") } if strings.ContainsAny(sessionID, `*?[\`+string(filepath.Separator)+"\x00") { return "", fmt.Errorf("pi: session id contains unsafe path characters") } stateDir := strings.TrimSpace(cfg.RuntimeStateDir) if stateDir == "" { stateDir = strings.TrimSpace(cfg.WorkDir) } if stateDir == "" { return "", fmt.Errorf("pi: runtime state dir is required") } dir := filepath.Join(stateDir, ".pi-sessions") matches, err := filepath.Glob(filepath.Join(dir, "*_"+sessionID+".jsonl")) if err != nil { return "", fmt.Errorf("pi: locate session %q: %w", sessionID, err) } if len(matches) == 0 { return "", fmt.Errorf("%w: %q", errPiSessionNotFound, sessionID) } for _, match := range matches { info, err := os.Lstat(match) if err != nil { return "", fmt.Errorf("pi: inspect session %q: %w", sessionID, err) } if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() { return "", fmt.Errorf("pi: session path %q is not a safe regular file", match) } } if len(matches) > 1 { best := matches[0] var bestMod time.Time for _, match := range matches { if info, err := os.Lstat(match); err == nil && info.ModTime().After(bestMod) { best = match bestMod = info.ModTime() } } return best, nil } return matches[0], nil }