11import { randomUUID } from "node:crypto" ;
2- import { unlinkSync , writeFileSync } from "node:fs" ;
2+ import { existsSync , unlinkSync , writeFileSync } from "node:fs" ;
33import { join } from "node:path" ;
44import type { AssistantMessage , Model , TextContent } from "@earendil-works/pi-ai" ;
55import {
6+ type AgentSessionEvent ,
67 AuthStorage ,
78 type CreateAgentSessionOptions ,
89 createAgentSession ,
@@ -21,6 +22,7 @@ import { classifyProviderLimit, WorkflowError, WorkflowErrorCode } from "./error
2122import { canonicalModelSpec , resolveModelSpecWithThinking } from "./model-spec.js" ;
2223import { loadModelTierConfig , type ModelTierConfig , resolveTierModel } from "./model-tier-config.js" ;
2324import { createStructuredOutputTool , type StructuredOutputCapture } from "./structured-output.js" ;
25+ import { workflowProjectPaths } from "./workflow-paths.js" ;
2426
2527/**
2628 * Find a JSON object/array in free-form text: a fenced ```json block if present,
@@ -215,9 +217,8 @@ export interface WorkflowAgentOptions {
215217 */
216218 modelRegistry ?: ModelRegistry ;
217219 /**
218- * Persist each subagent transcript as a real pi session file under the
219- * standard sessions directory (keyed by the runner's project cwd), instead
220- * of the default in-memory session that is discarded when the run ends.
220+ * Persist each subagent transcript as a private pi session file under the
221+ * workflow project's state directory, outside Pi's normal /resume picker.
221222 * Default: false (current behavior).
222223 */
223224 persistAgentSessions ?: boolean ;
@@ -243,14 +244,77 @@ export function listAvailableModelSpecs(registry?: ModelRegistry): string[] {
243244 }
244245}
245246
246- /** Real token /cost usage for a single subagent run, read from the SDK session . */
247+ /** Token /cost usage for a single subagent run. */
247248export interface AgentUsage {
248249 input : number ;
249250 output : number ;
250251 cacheRead : number ;
251252 cacheWrite : number ;
252253 total : number ;
253254 cost : number ;
255+ /** True only for an in-progress output-token estimate. */
256+ estimated ?: boolean ;
257+ }
258+
259+ /**
260+ * Convert session events into absolute cumulative usage. Exact message usage is
261+ * emitted at message_end; throttled message_update events add only a temporary
262+ * output estimate, which the next exact event replaces.
263+ */
264+ export function createAgentUsageEventHandler (
265+ onUsage : ( usage : AgentUsage ) => void ,
266+ now : ( ) => number = Date . now ,
267+ ) : ( event : AgentSessionEvent ) => void {
268+ const exact : AgentUsage = { input : 0 , output : 0 , cacheRead : 0 , cacheWrite : 0 , total : 0 , cost : 0 } ;
269+ const endedMessages = new WeakSet < object > ( ) ;
270+ let lastEstimateEmit = Number . NEGATIVE_INFINITY ;
271+ const emit = ( usage : AgentUsage ) => {
272+ try {
273+ onUsage ( usage ) ;
274+ } catch {
275+ // Telemetry is best-effort; never interrupt the child session.
276+ }
277+ } ;
278+
279+ return ( event ) => {
280+ if ( event . type === "message_end" && event . message . role === "assistant" ) {
281+ if ( endedMessages . has ( event . message ) ) return ;
282+ endedMessages . add ( event . message ) ;
283+ const usage = event . message . usage ;
284+ exact . input += usage . input ;
285+ exact . output += usage . output ;
286+ exact . cacheRead += usage . cacheRead ;
287+ exact . cacheWrite += usage . cacheWrite ;
288+ exact . total += usage . totalTokens ;
289+ exact . cost += usage . cost . total ;
290+ lastEstimateEmit = Number . NEGATIVE_INFINITY ;
291+ emit ( { ...exact , estimated : false } ) ;
292+ return ;
293+ }
294+
295+ if ( event . type !== "message_update" || event . message . role !== "assistant" ) return ;
296+ const timestamp = now ( ) ;
297+ if ( timestamp - lastEstimateEmit < 250 ) return ;
298+ lastEstimateEmit = timestamp ;
299+ const text = event . message . content
300+ . map ( ( part ) => ( part . type === "text" ? part . text : part . type === "thinking" ? part . thinking : "" ) )
301+ . join ( "" ) ;
302+ const estimatedOutput = Math . ceil ( text . length / 4 ) ;
303+ if ( estimatedOutput <= 0 ) return ;
304+ emit ( {
305+ ...exact ,
306+ output : exact . output + estimatedOutput ,
307+ total : exact . total + estimatedOutput ,
308+ estimated : true ,
309+ } ) ;
310+ } ;
311+ }
312+
313+ export interface AgentSessionCheckpoint {
314+ /** File-backed Pi child session. Undefined means this invocation cannot resume in-place. */
315+ sessionFile ?: string ;
316+ /** True when the existing session was reopened rather than created fresh. */
317+ resumed : boolean ;
254318}
255319
256320export interface AgentRunOptions < TSchemaDef extends TSchema | undefined = undefined > {
@@ -267,11 +331,15 @@ export interface AgentRunOptions<TSchemaDef extends TSchema | undefined = undefi
267331 instructions ?: string ;
268332 signal ?: AbortSignal ;
269333 /**
270- * Called once with this subagent's real usage, read from the session right
271- * before disposal. Fires on both the success and error paths so partial
272- * usage is never lost. `total === 0` means the provider reported no usage .
334+ * Called with absolute cumulative usage during the run and once more with
335+ * authoritative session totals before disposal. Streaming estimates have
336+ * `estimated: true`; message-boundary and terminal updates are exact .
273337 */
274338 onUsage ?: ( usage : AgentUsage ) => void ;
339+ /** Reopen this persisted Pi child session and continue from its last durable turn boundary. */
340+ resumeSessionFile ?: string ;
341+ /** Called before prompting so workflow persistence can durably link the invocation to its child session. */
342+ onSession ?: ( checkpoint : AgentSessionCheckpoint ) => void ;
275343 /**
276344 * Model spec for this subagent: either `provider/modelId` (unambiguous) or a
277345 * bare `modelId`. When it can't be resolved, the session default is used and
@@ -376,9 +444,9 @@ export class WorkflowAgent {
376444 }
377445
378446 /**
379- * Session manager for one subagent run. File-backed (persisted under the
380- * standard sessions dir, keyed by the runner's project cwd — never a
381- * per-call worktree cwd) when persistAgentSessions is on; in-memory otherwise.
447+ * Session manager for one subagent run. File-backed in the workflow project's
448+ * private agent- sessions directory (never the normal Pi /resume directory and
449+ * never a per-call worktree cwd) when persistence is on; in-memory otherwise.
382450 *
383451 * SessionManager.create() only creates the session directory — the SDK writes
384452 * the session file lazily (synchronous fs calls, uncaught) on the first
@@ -391,7 +459,7 @@ export class WorkflowAgent {
391459 private createSessionManager ( ) : SessionManager {
392460 if ( ! this . persistAgentSessions ) return SessionManager . inMemory ( ) ;
393461 try {
394- const manager = SessionManager . create ( this . cwd ) ;
462+ const manager = SessionManager . create ( this . cwd , workflowProjectPaths ( this . cwd ) . agentSessionsDir ) ;
395463 this . assertSessionDirWritable ( manager . getSessionDir ( ) ) ;
396464 return manager ;
397465 } catch ( error ) {
@@ -404,6 +472,28 @@ export class WorkflowAgent {
404472 }
405473 }
406474
475+ /** Reopen a durable child session when possible; otherwise create the configured fresh session. */
476+ private resolveSessionManager (
477+ resumeSessionFile : string | undefined ,
478+ runCwd : string ,
479+ ) : {
480+ manager : SessionManager ;
481+ resumed : boolean ;
482+ } {
483+ if ( resumeSessionFile && existsSync ( resumeSessionFile ) ) {
484+ try {
485+ return { manager : SessionManager . open ( resumeSessionFile , undefined , runCwd ) , resumed : true } ;
486+ } catch ( error ) {
487+ console . warn (
488+ `[workflow] could not reopen child session ${ resumeSessionFile } (${
489+ error instanceof Error ? error . message : String ( error )
490+ } ); restarting this agent from a fresh session`,
491+ ) ;
492+ }
493+ }
494+ return { manager : this . createSessionManager ( ) , resumed : false } ;
495+ }
496+
407497 /** Best-effort write probe: throws if the session directory isn't actually writable. */
408498 private assertSessionDirWritable ( dir : string ) : void {
409499 const probePath = join ( dir , `.write-probe-${ randomUUID ( ) } ` ) ;
@@ -463,11 +553,13 @@ export class WorkflowAgent {
463553 }
464554
465555 const agentDir = getAgentDir ( ) ;
466- // Key persisted sessions by the runner's project cwd (this.cwd), NOT the
467- // per-call runCwd: agents working in short-lived git worktrees should still
468- // group under the project's session dir instead of scattering across
469- // temporary worktree paths.
470- const sessionManager = this . createSessionManager ( ) ;
556+ // Key new persisted sessions by the runner's project cwd (this.cwd), NOT the
557+ // per-call runCwd. A resumed session is reopened with runCwd as its cwd override
558+ // so coding tools continue in the same shared tree or preserved worktree.
559+ const resolvedSession = this . sessionOptions . sessionManager
560+ ? { manager : this . sessionOptions . sessionManager , resumed : false }
561+ : this . resolveSessionManager ( options . resumeSessionFile , runCwd ) ;
562+ const sessionManager = resolvedSession . manager ;
471563 const { session } = await createAgentSession ( {
472564 cwd : runCwd ,
473565 agentDir,
@@ -489,18 +581,31 @@ export class WorkflowAgent {
489581 ...( resolvedThinkingLevel ? { thinkingLevel : resolvedThinkingLevel } : { } ) ,
490582 } ) ;
491583
492- // Name the persisted session so it's identifiable in session pickers.
493- // Skip when an injected session.sessionManager override won (tests/embedders).
494- if ( this . persistAgentSessions && ! this . sessionOptions . sessionManager && options . sessionName ) {
584+ // Persist the child-session link before the model starts, so a crash can recover it.
585+ try {
586+ options . onSession ?.( { sessionFile : sessionManager . getSessionFile ( ) , resumed : resolvedSession . resumed } ) ;
587+ } catch {
588+ // Session-link telemetry is best-effort; never block the child.
589+ }
590+
591+ // Name a newly persisted session so it's identifiable in session pickers.
592+ // Skip reopened sessions and injected managers to avoid duplicate session_info entries.
593+ if (
594+ this . persistAgentSessions &&
595+ ! resolvedSession . resumed &&
596+ ! this . sessionOptions . sessionManager &&
597+ options . sessionName
598+ ) {
495599 try {
496600 sessionManager . appendSessionInfo ( options . sessionName ) ;
497601 } catch {
498602 // Naming is best-effort; never fail the run over it.
499603 }
500604 }
501605
606+ const initialStats = resolvedSession . resumed ? session . getSessionStats ( ) : undefined ;
502607 let removeAbortListener : ( ( ) => void ) | undefined ;
503- let removeHistoryListener : ( ( ) => void ) | undefined ;
608+ let removeSessionListener : ( ( ) => void ) | undefined ;
504609 let lastHistoryEmit = 0 ;
505610 const emitHistory = ( ) => options . onHistory ?.( compactAgentHistory ( session . messages ) ) ;
506611 const maybeEmitHistory = ( ) => {
@@ -510,18 +615,26 @@ export class WorkflowAgent {
510615 lastHistoryEmit = now ;
511616 emitHistory ( ) ;
512617 } ;
618+ const handleUsageEvent = options . onUsage ? createAgentUsageEventHandler ( options . onUsage ) : undefined ;
513619 try {
514620 if ( options . signal ?. aborted ) throw new Error ( "Subagent was aborted" ) ;
515621 if ( options . signal ) {
516622 const onAbort = ( ) => void session . abort ( ) ;
517623 options . signal . addEventListener ( "abort" , onAbort , { once : true } ) ;
518624 removeAbortListener = ( ) => options . signal ?. removeEventListener ( "abort" , onAbort ) ;
519625 }
520- if ( options . onHistory ) {
521- removeHistoryListener = session . subscribe ( ( ) => maybeEmitHistory ( ) ) ;
626+ if ( options . onHistory || handleUsageEvent ) {
627+ removeSessionListener = session . subscribe ( ( event ) => {
628+ maybeEmitHistory ( ) ;
629+ handleUsageEvent ?.( event ) ;
630+ } ) ;
522631 }
523632
524- await session . prompt ( this . buildPrompt ( prompt , options as AgentRunOptions < any > , Boolean ( options . schema ) ) ) ;
633+ await session . prompt (
634+ resolvedSession . resumed
635+ ? this . buildResumePrompt ( Boolean ( options . schema ) )
636+ : this . buildPrompt ( prompt , options as AgentRunOptions < any > , Boolean ( options . schema ) ) ,
637+ ) ;
525638
526639 if ( options . signal ?. aborted ) throw new Error ( "Subagent was aborted" ) ;
527640
@@ -547,7 +660,7 @@ export class WorkflowAgent {
547660 return text as AgentRunResult < TSchemaDef > ;
548661 } finally {
549662 removeAbortListener ?.( ) ;
550- removeHistoryListener ?.( ) ;
663+ removeSessionListener ?.( ) ;
551664 try {
552665 emitHistory ( ) ;
553666 } catch {
@@ -558,12 +671,13 @@ export class WorkflowAgent {
558671 try {
559672 const { tokens, cost } = session . getSessionStats ( ) ;
560673 options . onUsage ( {
561- input : tokens . input ,
562- output : tokens . output ,
563- cacheRead : tokens . cacheRead ,
564- cacheWrite : tokens . cacheWrite ,
565- total : tokens . total ,
566- cost,
674+ input : Math . max ( 0 , tokens . input - ( initialStats ?. tokens . input ?? 0 ) ) ,
675+ output : Math . max ( 0 , tokens . output - ( initialStats ?. tokens . output ?? 0 ) ) ,
676+ cacheRead : Math . max ( 0 , tokens . cacheRead - ( initialStats ?. tokens . cacheRead ?? 0 ) ) ,
677+ cacheWrite : Math . max ( 0 , tokens . cacheWrite - ( initialStats ?. tokens . cacheWrite ?? 0 ) ) ,
678+ total : Math . max ( 0 , tokens . total - ( initialStats ?. tokens . total ?? 0 ) ) ,
679+ cost : Math . max ( 0 , cost - ( initialStats ?. cost ?? 0 ) ) ,
680+ estimated : false ,
567681 } ) ;
568682 } catch {
569683 // Usage is best-effort; never let stats failure mask the real result/error.
@@ -573,6 +687,20 @@ export class WorkflowAgent {
573687 }
574688 }
575689
690+ private buildResumePrompt ( structured : boolean ) : string {
691+ const parts = [
692+ "Continue the interrupted task from the existing session at the last durable message/tool-result boundary." ,
693+ "Review the prior messages and current filesystem state before acting. Do not repeat completed side effects." ,
694+ "Finish the original task and return its requested final result." ,
695+ ] ;
696+ if ( structured ) {
697+ parts . push (
698+ "When finished, call structured_output with the required schema, even if it was called before interruption." ,
699+ ) ;
700+ }
701+ return parts . join ( "\n\n" ) ;
702+ }
703+
576704 private buildPrompt ( prompt : string , options : AgentRunOptions < any > , structured : boolean ) : string {
577705 const parts = [
578706 this . instructions ,
0 commit comments