detached 4c138a0

fix: harden session transcript persistence

commit
4c138a07f871c362a4fe5cc10c65d1e5da25d42a
author
brice · · 2026-08-08 09:29 UTC
parent
3f37538
signature
GPG VerifiedCD88F9E564E0102646FB6077A1A07C7C7C3E349D · arch-gpg Verified for the committer email
10 files · +237 -78
M packages/agent/botcmd.go +2 -2
@@ -205,7 +205,7 @@ func botLogs(spec *botSpec, rawTail []string) error {
205205 }
206206
207207 // botRun starts the polling loop in the foreground. Ctrl+C stops it.
208-func botRun(spec *botSpec, rawTail []string, version string) error {
208+func botRun(spec *botSpec, rawTail []string, version string) (runErr error) {
209209 // Parse only a small subset of flags relevant to bot run. We reuse
210210 // the main args parser so --provider/--model/--cwd/--api-key/--reasoning
211211 // behave the same as in the tui.
@@ -252,7 +252,7 @@ func botRun(spec *botSpec, rawTail []string, version string) error {
252252 agent.OnTranscriptCompacted = func(msgs []provider.Message) {
253253 _ = sess.AppendCompaction(msgs)
254254 }
255- defer sess.Close()
255+ defer func() { runErr = joinSessionCloseError(runErr, sess) }()
256256 } else {
257257 fmt.Fprintln(os.Stderr, "session:", serr)
258258 }
@@ -205,7 +205,7 @@ func botLogs(spec *botSpec, rawTail []string) error {
205 }205 }
206 206
207 // botRun starts the polling loop in the foreground. Ctrl+C stops it.207 // botRun starts the polling loop in the foreground. Ctrl+C stops it.
208-func botRun(spec *botSpec, rawTail []string, version string) error {208+func botRun(spec *botSpec, rawTail []string, version string) (runErr error) {
209 // Parse only a small subset of flags relevant to bot run. We reuse209 // Parse only a small subset of flags relevant to bot run. We reuse
210 // the main args parser so --provider/--model/--cwd/--api-key/--reasoning210 // the main args parser so --provider/--model/--cwd/--api-key/--reasoning
211 // behave the same as in the tui.211 // behave the same as in the tui.
@@ -252,7 +252,7 @@ func botRun(spec *botSpec, rawTail []string, version string) error {
252 agent.OnTranscriptCompacted = func(msgs []provider.Message) {252 agent.OnTranscriptCompacted = func(msgs []provider.Message) {
253 _ = sess.AppendCompaction(msgs)253 _ = sess.AppendCompaction(msgs)
254 }254 }
255- defer sess.Close()255+ defer func() { runErr = joinSessionCloseError(runErr, sess) }()
256 } else {256 } else {
257 fmt.Fprintln(os.Stderr, "session:", serr)257 fmt.Fprintln(os.Stderr, "session:", serr)
258 }258 }
M packages/agent/cli.go +76 -44
@@ -351,6 +351,16 @@ func closeAgentLSP(ag *core.Agent) {
351351 }
352352 }
353353
354+func joinSessionCloseError(err error, sess *core.Session) error {
355+ if sess == nil {
356+ return err
357+ }
358+ if closeErr := sess.Close(); closeErr != nil {
359+ return errors.Join(err, fmt.Errorf("close session: %w", closeErr))
360+ }
361+ return err
362+}
363+
354364 func closeResolvedLSP(r Resolved) {
355365 _ = tools.CloseLSPManagers(r.ToolRegistry)
356366 }
@@ -443,7 +453,7 @@ func prepareSessionResumeWithOptions(path string, current *core.Agent, currentPr
443453 keepSession := false
444454 defer func() {
445455 if !keepSession {
446- _ = sess.Close()
456+ err = joinSessionCloseError(err, sess)
447457 }
448458 }()
449459
@@ -571,15 +581,16 @@ func applySessionResume(sess *core.Session, ag *core.Agent, currentProvider, cur
571581 }
572582 candidate, err := prepareSessionResumeWithOptions(sess.Path, ag, currentProvider, currentModel, explicitProvider, explicitModel, buildAgentFor)
573583 if err != nil {
574- _ = sess.Close()
575- return sessionResumeCandidate{}, err
584+ return sessionResumeCandidate{}, joinSessionCloseError(err, sess)
576585 }
577586 if !candidate.rebuilt {
578587 candidate.agent.SetMessages(candidate.messages)
579588 candidate.agent.SeedCost(candidate.cumulative)
580589 candidate.agent.SeedLastTurnUsage(candidate.lastTurn)
581590 }
582- _ = sess.Close()
591+ if closeErr := sess.Close(); closeErr != nil {
592+ return sessionResumeCandidate{}, errors.Join(fmt.Errorf("close session: %w", closeErr), joinSessionCloseError(nil, candidate.session))
593+ }
583594 return candidate, nil
584595 }
585596
@@ -879,7 +890,7 @@ func writePrintStats(path, providerID, model string, usage provider.Usage, elaps
879890 return nil
880891 }
881892
882-func runPrintMode(ctx context.Context, args Args, version string) error {
893+func runPrintMode(ctx context.Context, args Args, version string) (runErr error) {
883894 if args.NoYolo {
884895 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in print mode (no interactive prompt available); tools will run without confirmation")
885896 }
@@ -911,7 +922,7 @@ func runPrintMode(ctx context.Context, args Args, version string) error {
911922 }
912923 r.Provider, r.Model = providerName, model
913924 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }
914- defer sess.Close()
925+ defer func() { runErr = joinSessionCloseError(runErr, sess) }()
915926 }
916927 announceSession(extMgr, sess)
917928
@@ -957,7 +968,7 @@ func runPrintMode(ctx context.Context, args Args, version string) error {
957968 return nil
958969 }
959970
960-func runStreamMode(ctx context.Context, args Args, version string) error {
971+func runStreamMode(ctx context.Context, args Args, version string) (runErr error) {
961972 if args.NoYolo {
962973 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in stream mode (no interactive prompt available); tools will run without confirmation")
963974 }
@@ -989,7 +1000,7 @@ func runStreamMode(ctx context.Context, args Args, version string) error {
9891000 }
9901001 r.Provider, r.Model = providerName, model
9911002 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }
992- defer sess.Close()
1003+ defer func() { runErr = joinSessionCloseError(runErr, sess) }()
9931004 }
9941005 announceSession(extMgr, sess)
9951006
@@ -1074,7 +1085,7 @@ func newStreamTextSink(out io.Writer) (func(core.AgentEvent), func()) {
10741085 return sink, finish
10751086 }
10761087
1077-func runJSONMode(ctx context.Context, args Args, version string) error {
1088+func runJSONMode(ctx context.Context, args Args, version string) (runErr error) {
10781089 if args.NoYolo {
10791090 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in json mode (no interactive prompt available); tools will run without confirmation")
10801091 }
@@ -1106,7 +1117,7 @@ func runJSONMode(ctx context.Context, args Args, version string) error {
11061117 }
11071118 r.Provider, r.Model = providerName, model
11081119 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }
1109- defer sess.Close()
1120+ defer func() { runErr = joinSessionCloseError(runErr, sess) }()
11101121 }
11111122 announceSession(extMgr, sess)
11121123
@@ -1142,10 +1153,12 @@ func runJSONMode(ctx context.Context, args Args, version string) error {
11421153 transcriptStart = recovery.OutputStart
11431154 }
11441155 if persistErr := WriteNewTranscript(ag, sess, transcriptStart); persistErr != nil {
1156+ finalErr := persistErr
11451157 if err != nil {
1146- return errors.Join(err, persistErr)
1158+ finalErr = errors.Join(err, persistErr)
11471159 }
1148- return persistErr
1160+ _ = enc.Encode(map[string]any{"type": "error", "message": finalErr.Error()})
1161+ return finalErr
11491162 }
11501163 return err
11511164 }
@@ -1328,7 +1341,7 @@ func subagentPolicyFromConfig(cfg SubagentsConfig) subagents.SubagentPolicy {
13281341 }
13291342 }
13301343
1331-func runInteractive(ctx context.Context, args Args, version string) error {
1344+func runInteractive(ctx context.Context, args Args, version string) (runErr error) {
13321345 initialCfg, _ := LoadConfig()
13331346 // Resolve WITHOUT requiring credentials.
13341347 r, err := Resolve(args, false)
@@ -1762,9 +1775,7 @@ func runInteractive(ctx context.Context, args Args, version string) error {
17621775 defer func() {
17631776 persistMu.Lock()
17641777 defer persistMu.Unlock()
1765- if sess != nil {
1766- sess.Close()
1767- }
1778+ runErr = joinSessionCloseError(runErr, sess)
17681779 }()
17691780
17701781 // persistMessage is the per-message hook bound to the agent. It
@@ -1875,7 +1886,7 @@ func runInteractive(ctx context.Context, args Args, version string) error {
18751886
18761887 // loadSession replaces the current session with the one at path and
18771888 // hands its messages to the agent. Used by the /sessions picker.
1878- loadSession := func(path string) error {
1889+ loadSession := func(path string) (loadErr error) {
18791890 // Hold the transition lock from the pre-flush through the commit.
18801891 // Persistence callbacks take the read side, so an active session
18811892 // cannot be snapshotted before its lazy writes land or overwritten by
@@ -1896,9 +1907,11 @@ func runInteractive(ctx context.Context, args Args, version string) error {
18961907 // than the public callback, which would acquire the read lock held
18971908 // by this transition.
18981909 if oldSess != nil {
1899- if err := writeNewTranscriptLocked(currentAg, oldSess, sessBaselineMsgs); err != nil {
1910+ next, flushErr := writeNewTranscriptLocked(currentAg, oldSess, sessBaselineMsgs)
1911+ sessBaselineMsgs = next
1912+ if flushErr != nil {
19001913 persistMu.Unlock()
1901- return fmt.Errorf("flush current session: %w", err)
1914+ return fmt.Errorf("flush current session: %w", flushErr)
19021915 }
19031916 }
19041917 currentMessageCount := len(currentAg.Messages())
@@ -1916,7 +1929,7 @@ func runInteractive(ctx context.Context, args Args, version string) error {
19161929 committed := false
19171930 defer func() {
19181931 if !committed {
1919- _ = candidate.session.Close()
1932+ loadErr = joinSessionCloseError(loadErr, candidate.session)
19201933 }
19211934 }()
19221935
@@ -1940,8 +1953,9 @@ func runInteractive(ctx context.Context, args Args, version string) error {
19401953 return fmt.Errorf("session changed while loading; please try again")
19411954 }
19421955 newStates := copyExtensionStates(candidate.session.ExtensionState)
1956+ var oldCloseErr error
19431957 if oldSess != nil {
1944- _ = oldSess.Close()
1958+ oldCloseErr = oldSess.Close()
19451959 }
19461960 sess = candidate.session
19471961 sessionTitlePending = candidate.session != nil && candidate.session.Meta.Parent != "" && candidate.session.Title == "" && candidate.fullMessageCount <= candidate.session.Meta.ForkPoint
@@ -1982,6 +1996,9 @@ func runInteractive(ctx context.Context, args Args, version string) error {
19821996 extMgr.EmitSessionEvent("session_forked", sessionContext(candidate.session), newStates)
19831997 }
19841998 extMgr.EmitSessionEvent("session_switched", sessionContext(candidate.session), newStates)
1999+ if oldCloseErr != nil {
2000+ return fmt.Errorf("close current session: %w", oldCloseErr)
2001+ }
19852002 return nil
19862003 }
19872004
@@ -2132,15 +2149,18 @@ func runInteractive(ctx context.Context, args Args, version string) error {
21322149 // Commit only after every fallible operation has succeeded. Close
21332150 // the old session last so rollback never has to resurrect it.
21342151 var newStates map[string]json.RawMessage
2152+ var oldCloseErr error
21352153 persistMu.Lock()
21362154 if sess != nil {
2137- if err := writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs); err != nil {
2155+ next, flushErr := writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs)
2156+ sessBaselineMsgs = next
2157+ if flushErr != nil {
21382158 persistMu.Unlock()
2139- _ = newSess.Close()
2159+ rollbackErr := joinSessionCloseError(fmt.Errorf("flush current session: %w", flushErr), newSess)
21402160 rollback()
2141- return fmt.Errorf("flush current session: %w", err)
2161+ return rollbackErr
21422162 }
2143- _ = sess.Close()
2163+ oldCloseErr = sess.Close()
21442164 }
21452165 sess = newSess
21462166 if newSess != nil {
@@ -2171,6 +2191,9 @@ func runInteractive(ctx context.Context, args Args, version string) error {
21712191 if newSess != nil {
21722192 extMgr.EmitSessionEvent("session_opened", sessionContext(newSess), newStates)
21732193 }
2194+ if oldCloseErr != nil {
2195+ return fmt.Errorf("close current session: %w", oldCloseErr)
2196+ }
21742197 return nil
21752198 }
21762199
@@ -2447,11 +2470,12 @@ func runInteractive(ctx context.Context, args Args, version string) error {
24472470 if sess == nil {
24482471 return
24492472 }
2450- if err := writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs); err != nil {
2451- iv.ReportError(fmt.Errorf("flush session: %w", err))
2473+ next, flushErr := writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs)
2474+ sessBaselineMsgs = next
2475+ if flushErr != nil {
2476+ iv.ReportError(fmt.Errorf("flush session: %w", flushErr))
24522477 return
24532478 }
2454- sessBaselineMsgs = len(currentAg.Messages())
24552479 },
24562480 SessionTransition: newSessionTransition(&sessionTransitionMu),
24572481 Extensions: extMgr,
@@ -2570,12 +2594,14 @@ func runInteractive(ctx context.Context, args Args, version string) error {
25702594 if finalAg != nil {
25712595 persistMu.Lock()
25722596 if sess != nil {
2573- if err := writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs); err != nil {
2574- fmt.Fprintf(os.Stderr, "flush session: %v\n", err)
2575- } else {
2576- sessBaselineMsgs = len(finalAg.Messages())
2597+ next, flushErr := writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs)
2598+ sessBaselineMsgs = next
2599+ if flushErr != nil {
2600+ fmt.Fprintf(os.Stderr, "flush session: %v\n", flushErr)
2601+ }
2602+ if closeErr := sess.Close(); closeErr != nil {
2603+ fmt.Fprintf(os.Stderr, "close session: %v\n", closeErr)
25772604 }
2578- _ = sess.Close()
25792605 sess = nil
25802606 }
25812607 persistMu.Unlock()
@@ -2596,16 +2622,16 @@ func runInteractive(ctx context.Context, args Args, version string) error {
25962622 os.Exit(0)
25972623 }()
25982624
2599- runErr := iv.Run(ctx)
2625+ runErr = iv.Run(ctx)
26002626
26012627 // Flush final transcript to session (only if we had / ended up with an agent).
26022628 if finalAg := iv.Agent(); finalAg != nil {
26032629 persistMu.Lock()
26042630 if sess != nil {
2605- if err := writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs); err != nil {
2606- runErr = errors.Join(runErr, fmt.Errorf("flush session: %w", err))
2607- } else {
2608- sessBaselineMsgs = len(finalAg.Messages())
2631+ next, flushErr := writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs)
2632+ sessBaselineMsgs = next
2633+ if flushErr != nil {
2634+ runErr = errors.Join(runErr, fmt.Errorf("flush session: %w", flushErr))
26092635 }
26102636 }
26112637 persistMu.Unlock()
@@ -2704,28 +2730,34 @@ func pickSession(root, cwd string) (string, error) {
27042730 // the persistMu (non-interactive print/json modes which run a single
27052731 // turn under their own goroutine).
27062732 func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) error {
2707- return writeNewTranscriptLocked(ag, sess, from)
2733+ _, err := writeNewTranscriptLocked(ag, sess, from)
2734+ return err
27082735 }
27092736
27102737 // writeNewTranscriptLocked is the same as WriteNewTranscript. The
27112738 // suffix marks that interactive callers must hold persistMu when
27122739 // invoking it so concurrent appends from the agent loop don't race
27132740 // with this catch-up flush.
2714-func writeNewTranscriptLocked(ag *core.Agent, sess *core.Session, from int) error {
2741+func writeNewTranscriptLocked(ag *core.Agent, sess *core.Session, from int) (next int, err error) {
2742+ next = from
27152743 if sess == nil || ag == nil {
2716- return nil
2744+ return next, nil
27172745 }
27182746 msgs := ag.Messages()
27192747 for i := from; i < len(msgs); i++ {
27202748 if err := sess.AppendMessage(msgs[i]); err != nil {
2721- return fmt.Errorf("append transcript message %d: %w", i, err)
2749+ return next, fmt.Errorf("append transcript message %d: %w", i, err)
27222750 }
2751+ next = i + 1
27232752 }
27242753 cum := ag.Cost()
27252754 if err := sess.AppendUsage(cum, cum); err != nil {
2726- return fmt.Errorf("append transcript usage: %w", err)
2755+ return next, fmt.Errorf("append transcript usage: %w", err)
27272756 }
2728- return nil
2757+ if err := sess.Flush(); err != nil {
2758+ return next, fmt.Errorf("flush transcript: %w", err)
2759+ }
2760+ return next, nil
27292761 }
27302762
27312763 func readAllStdin() (string, error) {
@@ -351,6 +351,16 @@ func closeAgentLSP(ag *core.Agent) {
351 }351 }
352 }352 }
353 353
354+func joinSessionCloseError(err error, sess *core.Session) error {
355+ if sess == nil {
356+ return err
357+ }
358+ if closeErr := sess.Close(); closeErr != nil {
359+ return errors.Join(err, fmt.Errorf("close session: %w", closeErr))
360+ }
361+ return err
362+}
363+
354 func closeResolvedLSP(r Resolved) {364 func closeResolvedLSP(r Resolved) {
355 _ = tools.CloseLSPManagers(r.ToolRegistry)365 _ = tools.CloseLSPManagers(r.ToolRegistry)
356 }366 }
@@ -443,7 +453,7 @@ func prepareSessionResumeWithOptions(path string, current *core.Agent, currentPr
443 keepSession := false453 keepSession := false
444 defer func() {454 defer func() {
445 if !keepSession {455 if !keepSession {
446- _ = sess.Close()456+ err = joinSessionCloseError(err, sess)
447 }457 }
448 }()458 }()
449 459
@@ -571,15 +581,16 @@ func applySessionResume(sess *core.Session, ag *core.Agent, currentProvider, cur
571 }581 }
572 candidate, err := prepareSessionResumeWithOptions(sess.Path, ag, currentProvider, currentModel, explicitProvider, explicitModel, buildAgentFor)582 candidate, err := prepareSessionResumeWithOptions(sess.Path, ag, currentProvider, currentModel, explicitProvider, explicitModel, buildAgentFor)
573 if err != nil {583 if err != nil {
574- _ = sess.Close()584+ return sessionResumeCandidate{}, joinSessionCloseError(err, sess)
575- return sessionResumeCandidate{}, err
576 }585 }
577 if !candidate.rebuilt {586 if !candidate.rebuilt {
578 candidate.agent.SetMessages(candidate.messages)587 candidate.agent.SetMessages(candidate.messages)
579 candidate.agent.SeedCost(candidate.cumulative)588 candidate.agent.SeedCost(candidate.cumulative)
580 candidate.agent.SeedLastTurnUsage(candidate.lastTurn)589 candidate.agent.SeedLastTurnUsage(candidate.lastTurn)
581 }590 }
582- _ = sess.Close()591+ if closeErr := sess.Close(); closeErr != nil {
592+ return sessionResumeCandidate{}, errors.Join(fmt.Errorf("close session: %w", closeErr), joinSessionCloseError(nil, candidate.session))
593+ }
583 return candidate, nil594 return candidate, nil
584 }595 }
585 596
@@ -879,7 +890,7 @@ func writePrintStats(path, providerID, model string, usage provider.Usage, elaps
879 return nil890 return nil
880 }891 }
881 892
882-func runPrintMode(ctx context.Context, args Args, version string) error {893+func runPrintMode(ctx context.Context, args Args, version string) (runErr error) {
883 if args.NoYolo {894 if args.NoYolo {
884 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in print mode (no interactive prompt available); tools will run without confirmation")895 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in print mode (no interactive prompt available); tools will run without confirmation")
885 }896 }
@@ -911,7 +922,7 @@ func runPrintMode(ctx context.Context, args Args, version string) error {
911 }922 }
912 r.Provider, r.Model = providerName, model923 r.Provider, r.Model = providerName, model
913 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }924 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }
914- defer sess.Close()925+ defer func() { runErr = joinSessionCloseError(runErr, sess) }()
915 }926 }
916 announceSession(extMgr, sess)927 announceSession(extMgr, sess)
917 928
@@ -957,7 +968,7 @@ func runPrintMode(ctx context.Context, args Args, version string) error {
957 return nil968 return nil
958 }969 }
959 970
960-func runStreamMode(ctx context.Context, args Args, version string) error {971+func runStreamMode(ctx context.Context, args Args, version string) (runErr error) {
961 if args.NoYolo {972 if args.NoYolo {
962 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in stream mode (no interactive prompt available); tools will run without confirmation")973 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in stream mode (no interactive prompt available); tools will run without confirmation")
963 }974 }
@@ -989,7 +1000,7 @@ func runStreamMode(ctx context.Context, args Args, version string) error {
989 }1000 }
990 r.Provider, r.Model = providerName, model1001 r.Provider, r.Model = providerName, model
991 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }1002 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }
992- defer sess.Close()1003+ defer func() { runErr = joinSessionCloseError(runErr, sess) }()
993 }1004 }
994 announceSession(extMgr, sess)1005 announceSession(extMgr, sess)
995 1006
@@ -1074,7 +1085,7 @@ func newStreamTextSink(out io.Writer) (func(core.AgentEvent), func()) {
1074 return sink, finish1085 return sink, finish
1075 }1086 }
1076 1087
1077-func runJSONMode(ctx context.Context, args Args, version string) error {1088+func runJSONMode(ctx context.Context, args Args, version string) (runErr error) {
1078 if args.NoYolo {1089 if args.NoYolo {
1079 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in json mode (no interactive prompt available); tools will run without confirmation")1090 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in json mode (no interactive prompt available); tools will run without confirmation")
1080 }1091 }
@@ -1106,7 +1117,7 @@ func runJSONMode(ctx context.Context, args Args, version string) error {
1106 }1117 }
1107 r.Provider, r.Model = providerName, model1118 r.Provider, r.Model = providerName, model
1108 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }1119 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }
1109- defer sess.Close()1120+ defer func() { runErr = joinSessionCloseError(runErr, sess) }()
1110 }1121 }
1111 announceSession(extMgr, sess)1122 announceSession(extMgr, sess)
1112 1123
@@ -1142,10 +1153,12 @@ func runJSONMode(ctx context.Context, args Args, version string) error {
1142 transcriptStart = recovery.OutputStart1153 transcriptStart = recovery.OutputStart
1143 }1154 }
1144 if persistErr := WriteNewTranscript(ag, sess, transcriptStart); persistErr != nil {1155 if persistErr := WriteNewTranscript(ag, sess, transcriptStart); persistErr != nil {
1156+ finalErr := persistErr
1145 if err != nil {1157 if err != nil {
1146- return errors.Join(err, persistErr)1158+ finalErr = errors.Join(err, persistErr)
1147 }1159 }
1148- return persistErr1160+ _ = enc.Encode(map[string]any{"type": "error", "message": finalErr.Error()})
1161+ return finalErr
1149 }1162 }
1150 return err1163 return err
1151 }1164 }
@@ -1328,7 +1341,7 @@ func subagentPolicyFromConfig(cfg SubagentsConfig) subagents.SubagentPolicy {
1328 }1341 }
1329 }1342 }
1330 1343
1331-func runInteractive(ctx context.Context, args Args, version string) error {1344+func runInteractive(ctx context.Context, args Args, version string) (runErr error) {
1332 initialCfg, _ := LoadConfig()1345 initialCfg, _ := LoadConfig()
1333 // Resolve WITHOUT requiring credentials.1346 // Resolve WITHOUT requiring credentials.
1334 r, err := Resolve(args, false)1347 r, err := Resolve(args, false)
@@ -1762,9 +1775,7 @@ func runInteractive(ctx context.Context, args Args, version string) error {
1762 defer func() {1775 defer func() {
1763 persistMu.Lock()1776 persistMu.Lock()
1764 defer persistMu.Unlock()1777 defer persistMu.Unlock()
1765- if sess != nil {1778+ runErr = joinSessionCloseError(runErr, sess)
1766- sess.Close()
1767- }
1768 }()1779 }()
1769 1780
1770 // persistMessage is the per-message hook bound to the agent. It1781 // persistMessage is the per-message hook bound to the agent. It
@@ -1875,7 +1886,7 @@ func runInteractive(ctx context.Context, args Args, version string) error {
1875 1886
1876 // loadSession replaces the current session with the one at path and1887 // loadSession replaces the current session with the one at path and
1877 // hands its messages to the agent. Used by the /sessions picker.1888 // hands its messages to the agent. Used by the /sessions picker.
1878- loadSession := func(path string) error {1889+ loadSession := func(path string) (loadErr error) {
1879 // Hold the transition lock from the pre-flush through the commit.1890 // Hold the transition lock from the pre-flush through the commit.
1880 // Persistence callbacks take the read side, so an active session1891 // Persistence callbacks take the read side, so an active session
1881 // cannot be snapshotted before its lazy writes land or overwritten by1892 // cannot be snapshotted before its lazy writes land or overwritten by
@@ -1896,9 +1907,11 @@ func runInteractive(ctx context.Context, args Args, version string) error {
1896 // than the public callback, which would acquire the read lock held1907 // than the public callback, which would acquire the read lock held
1897 // by this transition.1908 // by this transition.
1898 if oldSess != nil {1909 if oldSess != nil {
1899- if err := writeNewTranscriptLocked(currentAg, oldSess, sessBaselineMsgs); err != nil {1910+ next, flushErr := writeNewTranscriptLocked(currentAg, oldSess, sessBaselineMsgs)
1911+ sessBaselineMsgs = next
1912+ if flushErr != nil {
1900 persistMu.Unlock()1913 persistMu.Unlock()
1901- return fmt.Errorf("flush current session: %w", err)1914+ return fmt.Errorf("flush current session: %w", flushErr)
1902 }1915 }
1903 }1916 }
1904 currentMessageCount := len(currentAg.Messages())1917 currentMessageCount := len(currentAg.Messages())
@@ -1916,7 +1929,7 @@ func runInteractive(ctx context.Context, args Args, version string) error {
1916 committed := false1929 committed := false
1917 defer func() {1930 defer func() {
1918 if !committed {1931 if !committed {
1919- _ = candidate.session.Close()1932+ loadErr = joinSessionCloseError(loadErr, candidate.session)
1920 }1933 }
1921 }()1934 }()
1922 1935
@@ -1940,8 +1953,9 @@ func runInteractive(ctx context.Context, args Args, version string) error {
1940 return fmt.Errorf("session changed while loading; please try again")1953 return fmt.Errorf("session changed while loading; please try again")
1941 }1954 }
1942 newStates := copyExtensionStates(candidate.session.ExtensionState)1955 newStates := copyExtensionStates(candidate.session.ExtensionState)
1956+ var oldCloseErr error
1943 if oldSess != nil {1957 if oldSess != nil {
1944- _ = oldSess.Close()1958+ oldCloseErr = oldSess.Close()
1945 }1959 }
1946 sess = candidate.session1960 sess = candidate.session
1947 sessionTitlePending = candidate.session != nil && candidate.session.Meta.Parent != "" && candidate.session.Title == "" && candidate.fullMessageCount <= candidate.session.Meta.ForkPoint1961 sessionTitlePending = candidate.session != nil && candidate.session.Meta.Parent != "" && candidate.session.Title == "" && candidate.fullMessageCount <= candidate.session.Meta.ForkPoint
@@ -1982,6 +1996,9 @@ func runInteractive(ctx context.Context, args Args, version string) error {
1982 extMgr.EmitSessionEvent("session_forked", sessionContext(candidate.session), newStates)1996 extMgr.EmitSessionEvent("session_forked", sessionContext(candidate.session), newStates)
1983 }1997 }
1984 extMgr.EmitSessionEvent("session_switched", sessionContext(candidate.session), newStates)1998 extMgr.EmitSessionEvent("session_switched", sessionContext(candidate.session), newStates)
1999+ if oldCloseErr != nil {
2000+ return fmt.Errorf("close current session: %w", oldCloseErr)
2001+ }
1985 return nil2002 return nil
1986 }2003 }
1987 2004
@@ -2132,15 +2149,18 @@ func runInteractive(ctx context.Context, args Args, version string) error {
2132 // Commit only after every fallible operation has succeeded. Close2149 // Commit only after every fallible operation has succeeded. Close
2133 // the old session last so rollback never has to resurrect it.2150 // the old session last so rollback never has to resurrect it.
2134 var newStates map[string]json.RawMessage2151 var newStates map[string]json.RawMessage
2152+ var oldCloseErr error
2135 persistMu.Lock()2153 persistMu.Lock()
2136 if sess != nil {2154 if sess != nil {
2137- if err := writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs); err != nil {2155+ next, flushErr := writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs)
2156+ sessBaselineMsgs = next
2157+ if flushErr != nil {
2138 persistMu.Unlock()2158 persistMu.Unlock()
2139- _ = newSess.Close()2159+ rollbackErr := joinSessionCloseError(fmt.Errorf("flush current session: %w", flushErr), newSess)
2140 rollback()2160 rollback()
2141- return fmt.Errorf("flush current session: %w", err)2161+ return rollbackErr
2142 }2162 }
2143- _ = sess.Close()2163+ oldCloseErr = sess.Close()
2144 }2164 }
2145 sess = newSess2165 sess = newSess
2146 if newSess != nil {2166 if newSess != nil {
@@ -2171,6 +2191,9 @@ func runInteractive(ctx context.Context, args Args, version string) error {
2171 if newSess != nil {2191 if newSess != nil {
2172 extMgr.EmitSessionEvent("session_opened", sessionContext(newSess), newStates)2192 extMgr.EmitSessionEvent("session_opened", sessionContext(newSess), newStates)
2173 }2193 }
2194+ if oldCloseErr != nil {
2195+ return fmt.Errorf("close current session: %w", oldCloseErr)
2196+ }
2174 return nil2197 return nil
2175 }2198 }
2176 2199
@@ -2447,11 +2470,12 @@ func runInteractive(ctx context.Context, args Args, version string) error {
2447 if sess == nil {2470 if sess == nil {
2448 return2471 return
2449 }2472 }
2450- if err := writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs); err != nil {2473+ next, flushErr := writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs)
2451- iv.ReportError(fmt.Errorf("flush session: %w", err))2474+ sessBaselineMsgs = next
2475+ if flushErr != nil {
2476+ iv.ReportError(fmt.Errorf("flush session: %w", flushErr))
2452 return2477 return
2453 }2478 }
2454- sessBaselineMsgs = len(currentAg.Messages())
2455 },2479 },
2456 SessionTransition: newSessionTransition(&sessionTransitionMu),2480 SessionTransition: newSessionTransition(&sessionTransitionMu),
2457 Extensions: extMgr,2481 Extensions: extMgr,
@@ -2570,12 +2594,14 @@ func runInteractive(ctx context.Context, args Args, version string) error {
2570 if finalAg != nil {2594 if finalAg != nil {
2571 persistMu.Lock()2595 persistMu.Lock()
2572 if sess != nil {2596 if sess != nil {
2573- if err := writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs); err != nil {2597+ next, flushErr := writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs)
2574- fmt.Fprintf(os.Stderr, "flush session: %v\n", err)2598+ sessBaselineMsgs = next
2575- } else {2599+ if flushErr != nil {
2576- sessBaselineMsgs = len(finalAg.Messages())2600+ fmt.Fprintf(os.Stderr, "flush session: %v\n", flushErr)
2601+ }
2602+ if closeErr := sess.Close(); closeErr != nil {
2603+ fmt.Fprintf(os.Stderr, "close session: %v\n", closeErr)
2577 }2604 }
2578- _ = sess.Close()
2579 sess = nil2605 sess = nil
2580 }2606 }
2581 persistMu.Unlock()2607 persistMu.Unlock()
@@ -2596,16 +2622,16 @@ func runInteractive(ctx context.Context, args Args, version string) error {
2596 os.Exit(0)2622 os.Exit(0)
2597 }()2623 }()
2598 2624
2599- runErr := iv.Run(ctx)2625+ runErr = iv.Run(ctx)
2600 2626
2601 // Flush final transcript to session (only if we had / ended up with an agent).2627 // Flush final transcript to session (only if we had / ended up with an agent).
2602 if finalAg := iv.Agent(); finalAg != nil {2628 if finalAg := iv.Agent(); finalAg != nil {
2603 persistMu.Lock()2629 persistMu.Lock()
2604 if sess != nil {2630 if sess != nil {
2605- if err := writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs); err != nil {2631+ next, flushErr := writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs)
2606- runErr = errors.Join(runErr, fmt.Errorf("flush session: %w", err))2632+ sessBaselineMsgs = next
2607- } else {2633+ if flushErr != nil {
2608- sessBaselineMsgs = len(finalAg.Messages())2634+ runErr = errors.Join(runErr, fmt.Errorf("flush session: %w", flushErr))
2609 }2635 }
2610 }2636 }
2611 persistMu.Unlock()2637 persistMu.Unlock()
@@ -2704,28 +2730,34 @@ func pickSession(root, cwd string) (string, error) {
2704 // the persistMu (non-interactive print/json modes which run a single2730 // the persistMu (non-interactive print/json modes which run a single
2705 // turn under their own goroutine).2731 // turn under their own goroutine).
2706 func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) error {2732 func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) error {
2707- return writeNewTranscriptLocked(ag, sess, from)2733+ _, err := writeNewTranscriptLocked(ag, sess, from)
2734+ return err
2708 }2735 }
2709 2736
2710 // writeNewTranscriptLocked is the same as WriteNewTranscript. The2737 // writeNewTranscriptLocked is the same as WriteNewTranscript. The
2711 // suffix marks that interactive callers must hold persistMu when2738 // suffix marks that interactive callers must hold persistMu when
2712 // invoking it so concurrent appends from the agent loop don't race2739 // invoking it so concurrent appends from the agent loop don't race
2713 // with this catch-up flush.2740 // with this catch-up flush.
2714-func writeNewTranscriptLocked(ag *core.Agent, sess *core.Session, from int) error {2741+func writeNewTranscriptLocked(ag *core.Agent, sess *core.Session, from int) (next int, err error) {
2742+ next = from
2715 if sess == nil || ag == nil {2743 if sess == nil || ag == nil {
2716- return nil2744+ return next, nil
2717 }2745 }
2718 msgs := ag.Messages()2746 msgs := ag.Messages()
2719 for i := from; i < len(msgs); i++ {2747 for i := from; i < len(msgs); i++ {
2720 if err := sess.AppendMessage(msgs[i]); err != nil {2748 if err := sess.AppendMessage(msgs[i]); err != nil {
2721- return fmt.Errorf("append transcript message %d: %w", i, err)2749+ return next, fmt.Errorf("append transcript message %d: %w", i, err)
2722 }2750 }
2751+ next = i + 1
2723 }2752 }
2724 cum := ag.Cost()2753 cum := ag.Cost()
2725 if err := sess.AppendUsage(cum, cum); err != nil {2754 if err := sess.AppendUsage(cum, cum); err != nil {
2726- return fmt.Errorf("append transcript usage: %w", err)2755+ return next, fmt.Errorf("append transcript usage: %w", err)
2727 }2756 }
2728- return nil2757+ if err := sess.Flush(); err != nil {
2758+ return next, fmt.Errorf("flush transcript: %w", err)
2759+ }
2760+ return next, nil
2729 }2761 }
2730 2762
2731 func readAllStdin() (string, error) {2763 func readAllStdin() (string, error) {
M packages/agent/cli_session_test.go +60 -0
@@ -347,3 +347,63 @@ func TestPrepareSessionResumeHonorsExplicitProviderModelFields(t *testing.T) {
347347 })
348348 }
349349 }
350+
351+func TestWriteNewTranscriptTracksPartialAppend(t *testing.T) {
352+ sess, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
353+ if err != nil {
354+ t.Fatal(err)
355+ }
356+ first := provider.Message{
357+ Role: provider.RoleUser,
358+ Content: []provider.Content{provider.TextBlock{Text: "first"}},
359+ }
360+ invalid := provider.Message{
361+ Role: provider.RoleAssistant,
362+ Content: []provider.Content{provider.ToolCallBlock{
363+ ID: "call-1",
364+ Name: "tool",
365+ Arguments: json.RawMessage("{"),
366+ }},
367+ }
368+ ag := core.NewAgent(nil, "model", "", nil)
369+ ag.SetMessages([]provider.Message{first, invalid})
370+ next, err := writeNewTranscriptLocked(ag, sess, 0)
371+ if err == nil {
372+ t.Fatal("writeNewTranscriptLocked succeeded with invalid message")
373+ }
374+ if next != 1 {
375+ t.Fatalf("next persisted message index = %d, want 1", next)
376+ }
377+
378+ second := provider.Message{
379+ Role: provider.RoleAssistant,
380+ Content: []provider.Content{provider.ToolCallBlock{ID: "call-1", Name: "tool", Arguments: json.RawMessage(`{}`)}},
381+ }
382+ third := provider.Message{
383+ Role: provider.RoleTool,
384+ Content: []provider.Content{provider.ToolResultBlock{CallID: "call-1", Content: []provider.Content{provider.TextBlock{Text: "done"}}}},
385+ }
386+ ag.SetMessages([]provider.Message{first, second, third})
387+ next, err = writeNewTranscriptLocked(ag, sess, next)
388+ if err != nil {
389+ t.Fatalf("retry transcript: %v", err)
390+ }
391+ if next != 3 {
392+ t.Fatalf("next persisted message index after retry = %d, want 3", next)
393+ }
394+ if err := sess.Close(); err != nil {
395+ t.Fatal(err)
396+ }
397+
398+ reopened, messages, err := core.OpenSession(sess.Path)
399+ if err != nil {
400+ t.Fatalf("reopen session: %v", err)
401+ }
402+ t.Cleanup(func() { _ = reopened.Close() })
403+ if len(messages) != 3 {
404+ t.Fatalf("replayed message count = %d, want 3", len(messages))
405+ }
406+ if repaired := provider.RepairOrphanedToolResults(messages); len(repaired) != len(messages) {
407+ t.Fatalf("replayed transcript lost tool pairing: %#v", messages)
408+ }
409+}
@@ -347,3 +347,63 @@ func TestPrepareSessionResumeHonorsExplicitProviderModelFields(t *testing.T) {
347 })347 })
348 }348 }
349 }349 }
350+
351+func TestWriteNewTranscriptTracksPartialAppend(t *testing.T) {
352+ sess, err := core.NewSession(t.TempDir(), "cwd", "provider", "model", "test")
353+ if err != nil {
354+ t.Fatal(err)
355+ }
356+ first := provider.Message{
357+ Role: provider.RoleUser,
358+ Content: []provider.Content{provider.TextBlock{Text: "first"}},
359+ }
360+ invalid := provider.Message{
361+ Role: provider.RoleAssistant,
362+ Content: []provider.Content{provider.ToolCallBlock{
363+ ID: "call-1",
364+ Name: "tool",
365+ Arguments: json.RawMessage("{"),
366+ }},
367+ }
368+ ag := core.NewAgent(nil, "model", "", nil)
369+ ag.SetMessages([]provider.Message{first, invalid})
370+ next, err := writeNewTranscriptLocked(ag, sess, 0)
371+ if err == nil {
372+ t.Fatal("writeNewTranscriptLocked succeeded with invalid message")
373+ }
374+ if next != 1 {
375+ t.Fatalf("next persisted message index = %d, want 1", next)
376+ }
377+
378+ second := provider.Message{
379+ Role: provider.RoleAssistant,
380+ Content: []provider.Content{provider.ToolCallBlock{ID: "call-1", Name: "tool", Arguments: json.RawMessage(`{}`)}},
381+ }
382+ third := provider.Message{
383+ Role: provider.RoleTool,
384+ Content: []provider.Content{provider.ToolResultBlock{CallID: "call-1", Content: []provider.Content{provider.TextBlock{Text: "done"}}}},
385+ }
386+ ag.SetMessages([]provider.Message{first, second, third})
387+ next, err = writeNewTranscriptLocked(ag, sess, next)
388+ if err != nil {
389+ t.Fatalf("retry transcript: %v", err)
390+ }
391+ if next != 3 {
392+ t.Fatalf("next persisted message index after retry = %d, want 3", next)
393+ }
394+ if err := sess.Close(); err != nil {
395+ t.Fatal(err)
396+ }
397+
398+ reopened, messages, err := core.OpenSession(sess.Path)
399+ if err != nil {
400+ t.Fatalf("reopen session: %v", err)
401+ }
402+ t.Cleanup(func() { _ = reopened.Close() })
403+ if len(messages) != 3 {
404+ t.Fatalf("replayed message count = %d, want 3", len(messages))
405+ }
406+ if repaired := provider.RepairOrphanedToolResults(messages); len(repaired) != len(messages) {
407+ t.Fatalf("replayed transcript lost tool pairing: %#v", messages)
408+ }
409+}
M packages/agent/modes/interactive.go +4 -1
@@ -9532,7 +9532,10 @@ func (i *Interactive) applySessionTreeTarget(target sessionTreeTarget, turnNo in
95329532 i.setSessionTreeError("tree: read selection: " + err.Error())
95339533 return
95349534 }
9535- _ = sess.Close()
9535+ if err := sess.Close(); err != nil {
9536+ i.setSessionTreeError("tree: close selection: " + err.Error())
9537+ return
9538+ }
95369539 msgs = current
95379540 }
95389541 selection, err := sessionTreeSelection(msgs, target)
@@ -9532,7 +9532,10 @@ func (i *Interactive) applySessionTreeTarget(target sessionTreeTarget, turnNo in
9532 i.setSessionTreeError("tree: read selection: " + err.Error())9532 i.setSessionTreeError("tree: read selection: " + err.Error())
9533 return9533 return
9534 }9534 }
9535- _ = sess.Close()9535+ if err := sess.Close(); err != nil {
9536+ i.setSessionTreeError("tree: close selection: " + err.Error())
9537+ return
9538+ }
9536 msgs = current9539 msgs = current
9537 }9540 }
9538 selection, err := sessionTreeSelection(msgs, target)9541 selection, err := sessionTreeSelection(msgs, target)
M packages/agent/subagent_worker.go +9 -15
@@ -36,7 +36,7 @@ import (
3636 // production; tests use the stubchild binary under
3737 // packages/agent/subagents/testdata/cmd/stubchild instead of the real model
3838 // loop.
39-func runSubagentWorkerMode(ctx context.Context, args Args, version string) error {
39+func runSubagentWorkerMode(ctx context.Context, args Args, version string) (runErr error) {
4040 if args.SubagentWorker == "" {
4141 return fmt.Errorf("--subagent-worker requires a socket path")
4242 }
@@ -86,7 +86,7 @@ func runSubagentWorkerMode(ctx context.Context, args Args, version string) error
8686 }
8787 r.Provider, r.Model = providerName, model
8888 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }
89- defer sess.Close()
89+ defer func() { runErr = joinSessionCloseError(runErr, sess) }()
9090 }
9191 announceSession(extMgr, sess)
9292
@@ -218,26 +218,21 @@ func runSubagentWorkerMode(ctx context.Context, args Args, version string) error
218218 cost: ag.Cost(),
219219 lastTurn: ag.LastTurnUsage(),
220220 }
221- compacted := false
222221 var persistCompaction func([]provider.Message) error
223222 if sess != nil {
224223 // Buffer the checkpoint until the continued turn completes. The final
225224 // transcript and cumulative usage are then persisted in one JSONL row.
226225 persistCompaction = func([]provider.Message) error {
227- if err := c.Err(); err != nil {
228- return err
229- }
230- compacted = true
231- return nil
226+ return c.Err()
232227 }
233228 }
234- start, err := promptWithContextRecovery(c, ag, prompt, sink, modes.ContextRecoveryOptions{
229+ recovery, err := promptWithContextRecovery(c, ag, prompt, sink, modes.ContextRecoveryOptions{
235230 PersistCompaction: persistCompaction,
236231 ContextWindow: contextWindow,
237232 AutoCompactThreshold: cfg.AutoCompactThreshold,
238233 })
239- output := finalAssistantText(ag.Messages()[start:])
240- if persistErr := persistWorkerTranscript(c, ag, sess, start, compacted, baseline); persistErr != nil {
234+ output := finalAssistantText(ag.Messages()[recovery.OutputStart:])
235+ if persistErr := persistWorkerTranscript(c, ag, sess, recovery.OutputStart, recovery.Compacted, baseline); persistErr != nil {
241236 if err != nil {
242237 err = errors.Join(err, persistErr)
243238 } else {
@@ -520,10 +515,9 @@ func persistWorkerTranscript(ctx context.Context, ag *core.Agent, sess *core.Ses
520515
521516 // promptWithContextRecovery preserves the worker's existing helper boundary
522517 // while delegating proactive compaction and one-shot overflow recovery to
523-// modes. The returned index is the first message to include in the result.
524-func promptWithContextRecovery(ctx context.Context, ag *core.Agent, prompt string, sink func(core.AgentEvent), options modes.ContextRecoveryOptions) (outputStart int, err error) {
525- result, err := modes.PromptWithContextRecovery(ctx, ag, prompt, nil, sink, options)
526- return result.OutputStart, err
518+// modes. The result includes both the output index and compaction state.
519+func promptWithContextRecovery(ctx context.Context, ag *core.Agent, prompt string, sink func(core.AgentEvent), options modes.ContextRecoveryOptions) (modes.ContextRecoveryResult, error) {
520+ return modes.PromptWithContextRecovery(ctx, ag, prompt, nil, sink, options)
527521 }
528522
529523 func finalAssistantText(messages []provider.Message) string {
@@ -36,7 +36,7 @@ import (
36 // production; tests use the stubchild binary under36 // production; tests use the stubchild binary under
37 // packages/agent/subagents/testdata/cmd/stubchild instead of the real model37 // packages/agent/subagents/testdata/cmd/stubchild instead of the real model
38 // loop.38 // loop.
39-func runSubagentWorkerMode(ctx context.Context, args Args, version string) error {39+func runSubagentWorkerMode(ctx context.Context, args Args, version string) (runErr error) {
40 if args.SubagentWorker == "" {40 if args.SubagentWorker == "" {
41 return fmt.Errorf("--subagent-worker requires a socket path")41 return fmt.Errorf("--subagent-worker requires a socket path")
42 }42 }
@@ -86,7 +86,7 @@ func runSubagentWorkerMode(ctx context.Context, args Args, version string) error
86 }86 }
87 r.Provider, r.Model = providerName, model87 r.Provider, r.Model = providerName, model
88 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }88 ag.OnToolResult = func(_ string, result core.ToolResult) { persistExtensionToolResult(extMgr, sess, result) }
89- defer sess.Close()89+ defer func() { runErr = joinSessionCloseError(runErr, sess) }()
90 }90 }
91 announceSession(extMgr, sess)91 announceSession(extMgr, sess)
92 92
@@ -218,26 +218,21 @@ func runSubagentWorkerMode(ctx context.Context, args Args, version string) error
218 cost: ag.Cost(),218 cost: ag.Cost(),
219 lastTurn: ag.LastTurnUsage(),219 lastTurn: ag.LastTurnUsage(),
220 }220 }
221- compacted := false
222 var persistCompaction func([]provider.Message) error221 var persistCompaction func([]provider.Message) error
223 if sess != nil {222 if sess != nil {
224 // Buffer the checkpoint until the continued turn completes. The final223 // Buffer the checkpoint until the continued turn completes. The final
225 // transcript and cumulative usage are then persisted in one JSONL row.224 // transcript and cumulative usage are then persisted in one JSONL row.
226 persistCompaction = func([]provider.Message) error {225 persistCompaction = func([]provider.Message) error {
227- if err := c.Err(); err != nil {226+ return c.Err()
228- return err
229- }
230- compacted = true
231- return nil
232 }227 }
233 }228 }
234- start, err := promptWithContextRecovery(c, ag, prompt, sink, modes.ContextRecoveryOptions{229+ recovery, err := promptWithContextRecovery(c, ag, prompt, sink, modes.ContextRecoveryOptions{
235 PersistCompaction: persistCompaction,230 PersistCompaction: persistCompaction,
236 ContextWindow: contextWindow,231 ContextWindow: contextWindow,
237 AutoCompactThreshold: cfg.AutoCompactThreshold,232 AutoCompactThreshold: cfg.AutoCompactThreshold,
238 })233 })
239- output := finalAssistantText(ag.Messages()[start:])234+ output := finalAssistantText(ag.Messages()[recovery.OutputStart:])
240- if persistErr := persistWorkerTranscript(c, ag, sess, start, compacted, baseline); persistErr != nil {235+ if persistErr := persistWorkerTranscript(c, ag, sess, recovery.OutputStart, recovery.Compacted, baseline); persistErr != nil {
241 if err != nil {236 if err != nil {
242 err = errors.Join(err, persistErr)237 err = errors.Join(err, persistErr)
243 } else {238 } else {
@@ -520,10 +515,9 @@ func persistWorkerTranscript(ctx context.Context, ag *core.Agent, sess *core.Ses
520 515
521 // promptWithContextRecovery preserves the worker's existing helper boundary516 // promptWithContextRecovery preserves the worker's existing helper boundary
522 // while delegating proactive compaction and one-shot overflow recovery to517 // while delegating proactive compaction and one-shot overflow recovery to
523-// modes. The returned index is the first message to include in the result.518+// modes. The result includes both the output index and compaction state.
524-func promptWithContextRecovery(ctx context.Context, ag *core.Agent, prompt string, sink func(core.AgentEvent), options modes.ContextRecoveryOptions) (outputStart int, err error) {519+func promptWithContextRecovery(ctx context.Context, ag *core.Agent, prompt string, sink func(core.AgentEvent), options modes.ContextRecoveryOptions) (modes.ContextRecoveryResult, error) {
525- result, err := modes.PromptWithContextRecovery(ctx, ag, prompt, nil, sink, options)520+ return modes.PromptWithContextRecovery(ctx, ag, prompt, nil, sink, options)
526- return result.OutputStart, err
527 }521 }
528 522
529 func finalAssistantText(messages []provider.Message) string {523 func finalAssistantText(messages []provider.Message) string {
M packages/agent/subagent_worker_test.go +19 -12
@@ -120,11 +120,14 @@ func TestPromptWithContextRecoveryCompactsAndContinues(t *testing.T) {
120120 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
121121 })
122122
123- outputStart, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{})
123+ recovery, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{})
124124 if err != nil {
125125 t.Fatalf("promptWithContextRecovery returned %v", err)
126126 }
127- if got := finalAssistantText(agent.Messages()[outputStart:]); got != "recovered" {
127+ if !recovery.Compacted {
128+ t.Fatal("promptWithContextRecovery did not report compaction")
129+ }
130+ if got := finalAssistantText(agent.Messages()[recovery.OutputStart:]); got != "recovered" {
128131 t.Fatalf("final assistant text = %q, want recovered", got)
129132 }
130133
@@ -245,7 +248,7 @@ func TestPromptWithContextRecoveryRetainsPromptWhenPersistenceFails(t *testing.T
245248 agent := core.NewAgent(client, "test-model", "", nil)
246249 agent.SetMessages(history)
247250
248- outputStart, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{
251+ recovery, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{
249252 PersistCompaction: func([]provider.Message) error { return errors.New("persist failed") },
250253 })
251254 if err == nil || !strings.Contains(err.Error(), "persist compacted transcript") {
@@ -254,7 +257,7 @@ func TestPromptWithContextRecoveryRetainsPromptWhenPersistenceFails(t *testing.T
254257 if got := workerRequestMessageTextCount(agent.Messages(), "prior context"); got != 1 || workerRequestMessageTextCount(agent.Messages(), "finish the task") != 1 {
255258 t.Fatalf("agent transcript lost the pending task after persistence failure: %#v", agent.Messages())
256259 }
257- if err := WriteNewTranscript(agent, session, outputStart); err != nil {
260+ if err := WriteNewTranscript(agent, session, recovery.OutputStart); err != nil {
258261 t.Fatalf("persist transcript: %v", err)
259262 }
260263 if err := session.Close(); err != nil {
@@ -300,13 +303,13 @@ func TestPromptWithContextRecoveryPersistsCompaction(t *testing.T) {
300303
301304 agent := core.NewAgent(client, "test-model", "", nil)
302305 agent.SetMessages(history)
303- outputStart, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{
306+ recovery, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{
304307 PersistCompaction: session.AppendCompaction,
305308 })
306309 if err != nil {
307310 t.Fatalf("promptWithContextRecovery returned %v", err)
308311 }
309- if err := WriteNewTranscript(agent, session, outputStart); err != nil {
312+ if err := WriteNewTranscript(agent, session, recovery.OutputStart); err != nil {
310313 t.Fatalf("persist transcript: %v", err)
311314 }
312315 if err := session.Close(); err != nil {
@@ -368,14 +371,14 @@ func TestWorkerUsesConfiguredAutoCompactionThreshold(t *testing.T) {
368371 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "answer"}}},
369372 })
370373 agent.SeedLastTurnUsage(provider.Usage{CacheReadTokens: 75})
371- start, err := promptWithContextRecovery(context.Background(), agent, "finish exact task", nil, modes.ContextRecoveryOptions{
374+ recovery, err := promptWithContextRecovery(context.Background(), agent, "finish exact task", nil, modes.ContextRecoveryOptions{
372375 ContextWindow: 100,
373376 AutoCompactThreshold: cfg.AutoCompactThreshold,
374377 })
375378 if err != nil {
376379 t.Fatalf("promptWithContextRecovery: %v", err)
377380 }
378- if got := finalAssistantText(agent.Messages()[start:]); got != "continued same task" {
381+ if got := finalAssistantText(agent.Messages()[recovery.OutputStart:]); got != "continued same task" {
379382 t.Fatalf("continued output = %q", got)
380383 }
381384 client.mu.Lock()
@@ -410,6 +413,10 @@ func (c *subagentSummaryUsageClient) Stream(_ context.Context, req provider.Requ
410413 if strings.Contains(req.System, "context summarization assistant") {
411414 events <- provider.EventUsage{Usage: provider.Usage{InputTokens: 3, OutputTokens: 1, CostUSD: 0.03}}
412415 events <- provider.EventTextDelta{Delta: "worker summary"}
416+ events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
417+ Role: provider.RoleAssistant,
418+ Content: []provider.Content{provider.TextBlock{Text: "worker summary"}},
419+ }}
413420 return
414421 }
415422 events <- provider.EventUsage{Usage: provider.Usage{InputTokens: 5, OutputTokens: 2, CostUSD: 0.05}}
@@ -443,20 +450,20 @@ func TestPersistWorkerSummaryUsageReopensWithCumulativeUsage(t *testing.T) {
443450 agent.SetMessages(history)
444451 agent.SeedCost(baseline)
445452 agent.SeedLastTurnUsage(provider.Usage{CacheReadTokens: 90})
446- outputStart, err := promptWithContextRecovery(context.Background(), agent, "continue task", nil, modes.ContextRecoveryOptions{
453+ recovery, err := promptWithContextRecovery(context.Background(), agent, "continue task", nil, modes.ContextRecoveryOptions{
447454 ContextWindow: 100,
448455 })
449456 if err != nil {
450457 t.Fatalf("promptWithContextRecovery: %v", err)
451458 }
452- if outputStart != len(history) {
453- t.Fatalf("compaction output start = %d, want %d", outputStart, len(history))
459+ if recovery.OutputStart != len(history) || !recovery.Compacted {
460+ t.Fatalf("compaction recovery = %+v, want output start %d and compacted", recovery, len(history))
454461 }
455462 expected := provider.Usage{InputTokens: 18, OutputTokens: 5, CostUSD: 0.18}
456463 if !reflect.DeepEqual(agent.Cost(), expected) {
457464 t.Fatalf("worker cumulative usage = %#v, want %#v", agent.Cost(), expected)
458465 }
459- if err := persistWorkerTranscript(context.Background(), agent, session, outputStart, true, workerTranscriptState{messages: history, cost: baseline}); err != nil {
466+ if err := persistWorkerTranscript(context.Background(), agent, session, recovery.OutputStart, recovery.Compacted, workerTranscriptState{messages: history, cost: baseline}); err != nil {
460467 t.Fatalf("persistWorkerTranscript: %v", err)
461468 }
462469 if err := session.Close(); err != nil {
@@ -120,11 +120,14 @@ func TestPromptWithContextRecoveryCompactsAndContinues(t *testing.T) {
120 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},120 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
121 })121 })
122 122
123- outputStart, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{})123+ recovery, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{})
124 if err != nil {124 if err != nil {
125 t.Fatalf("promptWithContextRecovery returned %v", err)125 t.Fatalf("promptWithContextRecovery returned %v", err)
126 }126 }
127- if got := finalAssistantText(agent.Messages()[outputStart:]); got != "recovered" {127+ if !recovery.Compacted {
128+ t.Fatal("promptWithContextRecovery did not report compaction")
129+ }
130+ if got := finalAssistantText(agent.Messages()[recovery.OutputStart:]); got != "recovered" {
128 t.Fatalf("final assistant text = %q, want recovered", got)131 t.Fatalf("final assistant text = %q, want recovered", got)
129 }132 }
130 133
@@ -245,7 +248,7 @@ func TestPromptWithContextRecoveryRetainsPromptWhenPersistenceFails(t *testing.T
245 agent := core.NewAgent(client, "test-model", "", nil)248 agent := core.NewAgent(client, "test-model", "", nil)
246 agent.SetMessages(history)249 agent.SetMessages(history)
247 250
248- outputStart, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{251+ recovery, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{
249 PersistCompaction: func([]provider.Message) error { return errors.New("persist failed") },252 PersistCompaction: func([]provider.Message) error { return errors.New("persist failed") },
250 })253 })
251 if err == nil || !strings.Contains(err.Error(), "persist compacted transcript") {254 if err == nil || !strings.Contains(err.Error(), "persist compacted transcript") {
@@ -254,7 +257,7 @@ func TestPromptWithContextRecoveryRetainsPromptWhenPersistenceFails(t *testing.T
254 if got := workerRequestMessageTextCount(agent.Messages(), "prior context"); got != 1 || workerRequestMessageTextCount(agent.Messages(), "finish the task") != 1 {257 if got := workerRequestMessageTextCount(agent.Messages(), "prior context"); got != 1 || workerRequestMessageTextCount(agent.Messages(), "finish the task") != 1 {
255 t.Fatalf("agent transcript lost the pending task after persistence failure: %#v", agent.Messages())258 t.Fatalf("agent transcript lost the pending task after persistence failure: %#v", agent.Messages())
256 }259 }
257- if err := WriteNewTranscript(agent, session, outputStart); err != nil {260+ if err := WriteNewTranscript(agent, session, recovery.OutputStart); err != nil {
258 t.Fatalf("persist transcript: %v", err)261 t.Fatalf("persist transcript: %v", err)
259 }262 }
260 if err := session.Close(); err != nil {263 if err := session.Close(); err != nil {
@@ -300,13 +303,13 @@ func TestPromptWithContextRecoveryPersistsCompaction(t *testing.T) {
300 303
301 agent := core.NewAgent(client, "test-model", "", nil)304 agent := core.NewAgent(client, "test-model", "", nil)
302 agent.SetMessages(history)305 agent.SetMessages(history)
303- outputStart, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{306+ recovery, err := promptWithContextRecovery(context.Background(), agent, "finish the task", nil, modes.ContextRecoveryOptions{
304 PersistCompaction: session.AppendCompaction,307 PersistCompaction: session.AppendCompaction,
305 })308 })
306 if err != nil {309 if err != nil {
307 t.Fatalf("promptWithContextRecovery returned %v", err)310 t.Fatalf("promptWithContextRecovery returned %v", err)
308 }311 }
309- if err := WriteNewTranscript(agent, session, outputStart); err != nil {312+ if err := WriteNewTranscript(agent, session, recovery.OutputStart); err != nil {
310 t.Fatalf("persist transcript: %v", err)313 t.Fatalf("persist transcript: %v", err)
311 }314 }
312 if err := session.Close(); err != nil {315 if err := session.Close(); err != nil {
@@ -368,14 +371,14 @@ func TestWorkerUsesConfiguredAutoCompactionThreshold(t *testing.T) {
368 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "answer"}}},371 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "answer"}}},
369 })372 })
370 agent.SeedLastTurnUsage(provider.Usage{CacheReadTokens: 75})373 agent.SeedLastTurnUsage(provider.Usage{CacheReadTokens: 75})
371- start, err := promptWithContextRecovery(context.Background(), agent, "finish exact task", nil, modes.ContextRecoveryOptions{374+ recovery, err := promptWithContextRecovery(context.Background(), agent, "finish exact task", nil, modes.ContextRecoveryOptions{
372 ContextWindow: 100,375 ContextWindow: 100,
373 AutoCompactThreshold: cfg.AutoCompactThreshold,376 AutoCompactThreshold: cfg.AutoCompactThreshold,
374 })377 })
375 if err != nil {378 if err != nil {
376 t.Fatalf("promptWithContextRecovery: %v", err)379 t.Fatalf("promptWithContextRecovery: %v", err)
377 }380 }
378- if got := finalAssistantText(agent.Messages()[start:]); got != "continued same task" {381+ if got := finalAssistantText(agent.Messages()[recovery.OutputStart:]); got != "continued same task" {
379 t.Fatalf("continued output = %q", got)382 t.Fatalf("continued output = %q", got)
380 }383 }
381 client.mu.Lock()384 client.mu.Lock()
@@ -410,6 +413,10 @@ func (c *subagentSummaryUsageClient) Stream(_ context.Context, req provider.Requ
410 if strings.Contains(req.System, "context summarization assistant") {413 if strings.Contains(req.System, "context summarization assistant") {
411 events <- provider.EventUsage{Usage: provider.Usage{InputTokens: 3, OutputTokens: 1, CostUSD: 0.03}}414 events <- provider.EventUsage{Usage: provider.Usage{InputTokens: 3, OutputTokens: 1, CostUSD: 0.03}}
412 events <- provider.EventTextDelta{Delta: "worker summary"}415 events <- provider.EventTextDelta{Delta: "worker summary"}
416+ events <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
417+ Role: provider.RoleAssistant,
418+ Content: []provider.Content{provider.TextBlock{Text: "worker summary"}},
419+ }}
413 return420 return
414 }421 }
415 events <- provider.EventUsage{Usage: provider.Usage{InputTokens: 5, OutputTokens: 2, CostUSD: 0.05}}422 events <- provider.EventUsage{Usage: provider.Usage{InputTokens: 5, OutputTokens: 2, CostUSD: 0.05}}
@@ -443,20 +450,20 @@ func TestPersistWorkerSummaryUsageReopensWithCumulativeUsage(t *testing.T) {
443 agent.SetMessages(history)450 agent.SetMessages(history)
444 agent.SeedCost(baseline)451 agent.SeedCost(baseline)
445 agent.SeedLastTurnUsage(provider.Usage{CacheReadTokens: 90})452 agent.SeedLastTurnUsage(provider.Usage{CacheReadTokens: 90})
446- outputStart, err := promptWithContextRecovery(context.Background(), agent, "continue task", nil, modes.ContextRecoveryOptions{453+ recovery, err := promptWithContextRecovery(context.Background(), agent, "continue task", nil, modes.ContextRecoveryOptions{
447 ContextWindow: 100,454 ContextWindow: 100,
448 })455 })
449 if err != nil {456 if err != nil {
450 t.Fatalf("promptWithContextRecovery: %v", err)457 t.Fatalf("promptWithContextRecovery: %v", err)
451 }458 }
452- if outputStart != len(history) {459+ if recovery.OutputStart != len(history) || !recovery.Compacted {
453- t.Fatalf("compaction output start = %d, want %d", outputStart, len(history))460+ t.Fatalf("compaction recovery = %+v, want output start %d and compacted", recovery, len(history))
454 }461 }
455 expected := provider.Usage{InputTokens: 18, OutputTokens: 5, CostUSD: 0.18}462 expected := provider.Usage{InputTokens: 18, OutputTokens: 5, CostUSD: 0.18}
456 if !reflect.DeepEqual(agent.Cost(), expected) {463 if !reflect.DeepEqual(agent.Cost(), expected) {
457 t.Fatalf("worker cumulative usage = %#v, want %#v", agent.Cost(), expected)464 t.Fatalf("worker cumulative usage = %#v, want %#v", agent.Cost(), expected)
458 }465 }
459- if err := persistWorkerTranscript(context.Background(), agent, session, outputStart, true, workerTranscriptState{messages: history, cost: baseline}); err != nil {466+ if err := persistWorkerTranscript(context.Background(), agent, session, recovery.OutputStart, recovery.Compacted, workerTranscriptState{messages: history, cost: baseline}); err != nil {
460 t.Fatalf("persistWorkerTranscript: %v", err)467 t.Fatalf("persistWorkerTranscript: %v", err)
461 }468 }
462 if err := session.Close(); err != nil {469 if err := session.Close(); err != nil {
M packages/core/agent.go +7 -1
@@ -347,6 +347,12 @@ func (a *Agent) Cost() provider.Usage {
347347 return a.cost.Total
348348 }
349349
350+func (a *Agent) addUsage(usage provider.Usage) provider.Usage {
351+ a.mu.Lock()
352+ defer a.mu.Unlock()
353+ return a.cost.Add(usage)
354+}
355+
350356 // SeedCost sets the cumulative usage as a baseline before the first
351357 // turn runs. Used when transferring state from another agent (model
352358 // or provider switch) so the running cost meter doesn't reset to 0.
@@ -750,7 +756,7 @@ func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent), turnContext
750756 case provider.EventToolEnd:
751757 sink(EvToolUseEnd{ID: e.ID})
752758 case provider.EventUsage:
753- cum := a.cost.Add(e.Usage)
759+ cum := a.addUsage(e.Usage)
754760 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
755761 if a.OnUsage != nil {
756762 a.OnUsage(cum)
@@ -347,6 +347,12 @@ func (a *Agent) Cost() provider.Usage {
347 return a.cost.Total347 return a.cost.Total
348 }348 }
349 349
350+func (a *Agent) addUsage(usage provider.Usage) provider.Usage {
351+ a.mu.Lock()
352+ defer a.mu.Unlock()
353+ return a.cost.Add(usage)
354+}
355+
350 // SeedCost sets the cumulative usage as a baseline before the first356 // SeedCost sets the cumulative usage as a baseline before the first
351 // turn runs. Used when transferring state from another agent (model357 // turn runs. Used when transferring state from another agent (model
352 // or provider switch) so the running cost meter doesn't reset to 0.358 // or provider switch) so the running cost meter doesn't reset to 0.
@@ -750,7 +756,7 @@ func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent), turnContext
750 case provider.EventToolEnd:756 case provider.EventToolEnd:
751 sink(EvToolUseEnd{ID: e.ID})757 sink(EvToolUseEnd{ID: e.ID})
752 case provider.EventUsage:758 case provider.EventUsage:
753- cum := a.cost.Add(e.Usage)759+ cum := a.addUsage(e.Usage)
754 sink(EvUsage{Usage: e.Usage, Cumulative: cum})760 sink(EvUsage{Usage: e.Usage, Cumulative: cum})
755 if a.OnUsage != nil {761 if a.OnUsage != nil {
756 a.OnUsage(cum)762 a.OnUsage(cum)
M packages/core/compact.go +1 -1
@@ -111,7 +111,7 @@ func (a *Agent) compact(ctx context.Context, keepTail int, textSink func(delta s
111111 eventSink(EvTextDelta{Delta: e.Delta})
112112 }
113113 case provider.EventUsage:
114- cum := a.cost.Add(e.Usage)
114+ cum := a.addUsage(e.Usage)
115115 if eventSink != nil {
116116 eventSink(EvUsage{Usage: e.Usage, Cumulative: cum})
117117 }
@@ -111,7 +111,7 @@ func (a *Agent) compact(ctx context.Context, keepTail int, textSink func(delta s
111 eventSink(EvTextDelta{Delta: e.Delta})111 eventSink(EvTextDelta{Delta: e.Delta})
112 }112 }
113 case provider.EventUsage:113 case provider.EventUsage:
114- cum := a.cost.Add(e.Usage)114+ cum := a.addUsage(e.Usage)
115 if eventSink != nil {115 if eventSink != nil {
116 eventSink(EvUsage{Usage: e.Usage, Cumulative: cum})116 eventSink(EvUsage{Usage: e.Usage, Cumulative: cum})
117 }117 }
M packages/core/compact_lifecycle_test.go +50 -1
@@ -12,6 +12,7 @@ type compactLifecycleClient struct {
1212 hadLifecycle bool
1313 req provider.Request
1414 summary string
15+ usage []provider.Usage
1516 }
1617
1718 func (c *compactLifecycleClient) Name() string { return "compact-lifecycle" }
@@ -26,8 +27,11 @@ func (c *compactLifecycleClient) Stream(_ context.Context, req provider.Request)
2627 if summary == "" {
2728 summary = "summary"
2829 }
29- out := make(chan provider.Event, 2)
30+ out := make(chan provider.Event, 2+len(c.usage))
3031 out <- provider.EventTextDelta{Delta: summary}
32+ for _, usage := range c.usage {
33+ out <- provider.EventUsage{Usage: usage}
34+ }
3135 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
3236 Role: provider.RoleAssistant,
3337 Content: []provider.Content{provider.TextBlock{Text: summary}},
@@ -36,6 +40,51 @@ func (c *compactLifecycleClient) Stream(_ context.Context, req provider.Request)
3640 return out, nil
3741 }
3842
43+func TestCompactWithEventsAccumulatesAndPersistsUsage(t *testing.T) {
44+ client := &compactLifecycleClient{usage: []provider.Usage{{InputTokens: 2}, {OutputTokens: 3, CostUSD: 0.25}}}
45+ agent := NewAgent(client, "compact-model", "system", Registry{})
46+ agent.SetMessages([]provider.Message{{
47+ Role: provider.RoleUser,
48+ Content: []provider.Content{provider.TextBlock{Text: "source"}},
49+ }})
50+ root := t.TempDir()
51+ sess, err := NewSession(root, t.TempDir(), "provider", "compact-model", "test")
52+ if err != nil {
53+ t.Fatal(err)
54+ }
55+ t.Cleanup(func() { _ = sess.Close() })
56+ agent.OnUsage = func(cumulative provider.Usage) {
57+ if err := sess.AppendUsage(cumulative, cumulative); err != nil {
58+ t.Fatalf("persist usage: %v", err)
59+ }
60+ }
61+ var events []AgentEvent
62+ if _, err := agent.CompactWithEvents(context.Background(), 0, func(event AgentEvent) {
63+ events = append(events, event)
64+ }); err != nil {
65+ t.Fatalf("CompactWithEvents: %v", err)
66+ }
67+ if got, want := agent.Cost(), (provider.Usage{InputTokens: 2, OutputTokens: 3, CostUSD: 0.25}); got != want {
68+ t.Fatalf("agent cost = %+v, want %+v", got, want)
69+ }
70+ cum, _, err := SessionUsageDetail(sess.Path)
71+ if err != nil {
72+ t.Fatal(err)
73+ }
74+ if cum != agent.Cost() {
75+ t.Fatalf("persisted cumulative usage = %+v, want %+v", cum, agent.Cost())
76+ }
77+ var cumulative []provider.Usage
78+ for _, event := range events {
79+ if usage, ok := event.(EvUsage); ok {
80+ cumulative = append(cumulative, usage.Cumulative)
81+ }
82+ }
83+ if len(cumulative) != 2 || cumulative[0].InputTokens != 2 || cumulative[1] != agent.Cost() {
84+ t.Fatalf("usage event cumulative values = %+v, want first input=2 and final=%+v", cumulative, agent.Cost())
85+ }
86+}
87+
3988 func TestCompactWithEventsEmitsRequestLifecycle(t *testing.T) {
4089 client := &compactLifecycleClient{}
4190 agent := NewAgent(client, "compact-model", "system", Registry{})
@@ -12,6 +12,7 @@ type compactLifecycleClient struct {
12 hadLifecycle bool12 hadLifecycle bool
13 req provider.Request13 req provider.Request
14 summary string14 summary string
15+ usage []provider.Usage
15 }16 }
16 17
17 func (c *compactLifecycleClient) Name() string { return "compact-lifecycle" }18 func (c *compactLifecycleClient) Name() string { return "compact-lifecycle" }
@@ -26,8 +27,11 @@ func (c *compactLifecycleClient) Stream(_ context.Context, req provider.Request)
26 if summary == "" {27 if summary == "" {
27 summary = "summary"28 summary = "summary"
28 }29 }
29- out := make(chan provider.Event, 2)30+ out := make(chan provider.Event, 2+len(c.usage))
30 out <- provider.EventTextDelta{Delta: summary}31 out <- provider.EventTextDelta{Delta: summary}
32+ for _, usage := range c.usage {
33+ out <- provider.EventUsage{Usage: usage}
34+ }
31 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{35 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
32 Role: provider.RoleAssistant,36 Role: provider.RoleAssistant,
33 Content: []provider.Content{provider.TextBlock{Text: summary}},37 Content: []provider.Content{provider.TextBlock{Text: summary}},
@@ -36,6 +40,51 @@ func (c *compactLifecycleClient) Stream(_ context.Context, req provider.Request)
36 return out, nil40 return out, nil
37 }41 }
38 42
43+func TestCompactWithEventsAccumulatesAndPersistsUsage(t *testing.T) {
44+ client := &compactLifecycleClient{usage: []provider.Usage{{InputTokens: 2}, {OutputTokens: 3, CostUSD: 0.25}}}
45+ agent := NewAgent(client, "compact-model", "system", Registry{})
46+ agent.SetMessages([]provider.Message{{
47+ Role: provider.RoleUser,
48+ Content: []provider.Content{provider.TextBlock{Text: "source"}},
49+ }})
50+ root := t.TempDir()
51+ sess, err := NewSession(root, t.TempDir(), "provider", "compact-model", "test")
52+ if err != nil {
53+ t.Fatal(err)
54+ }
55+ t.Cleanup(func() { _ = sess.Close() })
56+ agent.OnUsage = func(cumulative provider.Usage) {
57+ if err := sess.AppendUsage(cumulative, cumulative); err != nil {
58+ t.Fatalf("persist usage: %v", err)
59+ }
60+ }
61+ var events []AgentEvent
62+ if _, err := agent.CompactWithEvents(context.Background(), 0, func(event AgentEvent) {
63+ events = append(events, event)
64+ }); err != nil {
65+ t.Fatalf("CompactWithEvents: %v", err)
66+ }
67+ if got, want := agent.Cost(), (provider.Usage{InputTokens: 2, OutputTokens: 3, CostUSD: 0.25}); got != want {
68+ t.Fatalf("agent cost = %+v, want %+v", got, want)
69+ }
70+ cum, _, err := SessionUsageDetail(sess.Path)
71+ if err != nil {
72+ t.Fatal(err)
73+ }
74+ if cum != agent.Cost() {
75+ t.Fatalf("persisted cumulative usage = %+v, want %+v", cum, agent.Cost())
76+ }
77+ var cumulative []provider.Usage
78+ for _, event := range events {
79+ if usage, ok := event.(EvUsage); ok {
80+ cumulative = append(cumulative, usage.Cumulative)
81+ }
82+ }
83+ if len(cumulative) != 2 || cumulative[0].InputTokens != 2 || cumulative[1] != agent.Cost() {
84+ t.Fatalf("usage event cumulative values = %+v, want first input=2 and final=%+v", cumulative, agent.Cost())
85+ }
86+}
87+
39 func TestCompactWithEventsEmitsRequestLifecycle(t *testing.T) {88 func TestCompactWithEventsEmitsRequestLifecycle(t *testing.T) {
40 client := &compactLifecycleClient{}89 client := &compactLifecycleClient{}
41 agent := NewAgent(client, "compact-model", "system", Registry{})90 agent := NewAgent(client, "compact-model", "system", Registry{})
M packages/core/session.go +9 -1
@@ -1327,6 +1327,14 @@ func (s *Session) AppendExtensionState(extension string, state json.RawMessage)
13271327 return nil
13281328 }
13291329
1330+// Flush writes buffered session data to the append handle.
1331+func (s *Session) Flush() error {
1332+ if s == nil {
1333+ return nil
1334+ }
1335+ return s.buf.Flush()
1336+}
1337+
13301338 // Close flushes and closes the session file. If the session was
13311339 // freshly created in this process and never had any messages
13321340 // appended (the user opened zut, looked around, and exited without
@@ -1336,7 +1344,7 @@ func (s *Session) Close() error {
13361344 if s == nil {
13371345 return nil
13381346 }
1339- flushErr := s.buf.Flush()
1347+ flushErr := s.Flush()
13401348 closeErr := s.writer.Close()
13411349 if s.freshFile && s.messagesAppended == 0 && len(s.ExtensionState) == 0 && len(s.Meta.CompactHandoff) == 0 {
13421350 // Best-effort cleanup. We deliberately don't propagate the
@@ -1327,6 +1327,14 @@ func (s *Session) AppendExtensionState(extension string, state json.RawMessage)
1327 return nil1327 return nil
1328 }1328 }
1329 1329
1330+// Flush writes buffered session data to the append handle.
1331+func (s *Session) Flush() error {
1332+ if s == nil {
1333+ return nil
1334+ }
1335+ return s.buf.Flush()
1336+}
1337+
1330 // Close flushes and closes the session file. If the session was1338 // Close flushes and closes the session file. If the session was
1331 // freshly created in this process and never had any messages1339 // freshly created in this process and never had any messages
1332 // appended (the user opened zut, looked around, and exited without1340 // appended (the user opened zut, looked around, and exited without
@@ -1336,7 +1344,7 @@ func (s *Session) Close() error {
1336 if s == nil {1344 if s == nil {
1337 return nil1345 return nil
1338 }1346 }
1339- flushErr := s.buf.Flush()1347+ flushErr := s.Flush()
1340 closeErr := s.writer.Close()1348 closeErr := s.writer.Close()
1341 if s.freshFile && s.messagesAppended == 0 && len(s.ExtensionState) == 0 && len(s.Meta.CompactHandoff) == 0 {1349 if s.freshFile && s.messagesAppended == 0 && len(s.ExtensionState) == 0 && len(s.Meta.CompactHandoff) == 0 {
1342 // Best-effort cleanup. We deliberately don't propagate the1350 // Best-effort cleanup. We deliberately don't propagate the

Keyboard shortcuts

Single-key shortcuts are on. ctrl k and esc work whatever this says.

ctrl k :
go to a section, branch, tag, file or repository
?
this list
j k
move the row cursor down / up
enter
open the row under the cursor; on Files, a directory opens or closes in place
↑ ↓ ← → home end
in the file tree (focus on a row): previous / next row, close / open a directory or go to its parent / first entry, first / last row. They never open a file
l h
on Files: open a directory or a file / close a directory or go up to its parent, then to the parent page. Elsewhere: into the side panel (on a merge request, its files) / back to the content; there j k move between its controls
g g G
first / last row
/
filter the list; on Files, find a file anywhere in the repository (esc clears, enter opens the best match)
b
switch branch or tag
g s g f g c g b g m g i g r g a
tabs: summary, files, commits, branches & tags, merge requests, issues, releases, actions (when available)
g h
home
w
wide mode: use the whole width of the screen, or not (remembered)
y
copy a permalink pinned to the commit
d u x
diff: split / unified / collapse or expand all
e E
diff: the whole file in view / every file, with the lines the hunks hide
n
merge request: next note in the diff
esc
close a dialog, leave a field

Branches 1

Tags 477

↑↓ moveenter openesc close