session: persist each turn as it happens, flush on SIGTERM/SIGHUP
Until now an interactive session's turns lived only in the running agent's in-memory transcript; the session file on disk got nothing new until iv.Run returned and the deferred WriteNewTranscript at the bottom of runInteractive fired. That meant any process death that bypassed the deferred flush \u2014 closed terminal window (SIGHUP), system shutdown (SIGTERM), kill -9, an OS / power crash \u2014 took the entire session with it. The window of data-at-risk was "everything since the TUI started or last switched session." Two changes that close the window: - Add OnMessageAppended / OnUsage hooks on core.Agent, fired right after each transcript message is appended (user prompt, finalised assistant turn, tool-results message, the OpenAI image mirror) and after each usage event arrives. The interactive runtime in cli.go binds them to sess.AppendMessage / sess.AppendUsage so every finished turn lands in the session JSONL the moment it's durable in memory. A persistMu mutex coordinates the agent goroutine's per-message writes with the TUI goroutine's session swap (/sessions) and explicit flushes (/session export). sessBaselineMsgs advances in lock-step so the exit-time WriteNewTranscript no longer double-writes rows already on disk. - Install a SIGTERM / SIGHUP handler in runInteractive that flushes any not-yet-persisted in-flight turn before exiting. SIGINT is intentionally NOT handled \u2014 the TUI consumes Ctrl+C as a regular key event for cancel/clear semantics, so installing a SIGINT notifier here would swallow it. The handler exits with os.Exit(0) rather than re-raising, because re-raising would skip the chance to flush and the only at-risk state we care about (the session file) is already flushed by the time we get there. Build closures (BuildAgent / BuildAgentFor) are re-wrapped after the persistence wiring exists so any agent the TUI constructs on login or /model swap also gets the hooks; otherwise switching provider would silently revert to the old in-memory-only behaviour. Print and JSON modes are unaffected: they run a single turn and already flush via WriteNewTranscript at the end of their handler.
- commit
- 65135596a0bd8b8994a348cbde03e2afd91a889c
- author
- patriceckhart · · 2026-04-28 06:01 UTC
- parent
- 853a089
- signature
- unsigned
- refs
- tag: v0.1.32
2 files · +168 -12
M internal/agent/cli.go +132 -12
| @@ -7,7 +7,10 @@ import ( | |||
| 7 | 7 | "fmt" | |
| 8 | 8 | "io" | |
| 9 | 9 | "os" | |
| 10 | + | "os/signal" | |
| 10 | 11 | "strings" | |
| 12 | + | "sync" | |
| 13 | + | "syscall" | |
| 11 | 14 | "time" | |
| 12 | 15 | ||
| 13 | 16 | "github.com/patriceckhart/zot/internal/agent/extensions" | |
| @@ -482,6 +485,12 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 482 | 485 | ||
| 483 | 486 | var sess *core.Session | |
| 484 | 487 | var sessBaselineMsgs int // messages already on disk when current session opened | |
| 488 | + | // persistMu guards sess + sessBaselineMsgs against concurrent access | |
| 489 | + | // from the agent loop's per-message persistence hook (runs on the | |
| 490 | + | // agent goroutine) and the TUI's session swap / flush callbacks | |
| 491 | + | // (run on the TUI goroutine). Without this, a /sessions swap that | |
| 492 | + | // races with a finishing turn could double-write or lose messages. | |
| 493 | + | var persistMu sync.Mutex | |
| 485 | 494 | if !args.NoSess && ag != nil { | |
| 486 | 495 | sess, _ = openOrCreateSession(args, r, ag, version) | |
| 487 | 496 | if ag != nil { | |
| @@ -489,11 +498,62 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 489 | 498 | } | |
| 490 | 499 | } | |
| 491 | 500 | defer func() { | |
| 501 | + | persistMu.Lock() | |
| 502 | + | defer persistMu.Unlock() | |
| 492 | 503 | if sess != nil { | |
| 493 | 504 | sess.Close() | |
| 494 | 505 | } | |
| 495 | 506 | }() | |
| 496 | 507 | ||
| 508 | + | // persistMessage is the per-message hook bound to the agent. It | |
| 509 | + | // appends each new transcript message to the live session as soon | |
| 510 | + | // as it lands, so a kill / closed terminal / OS crash costs at | |
| 511 | + | // most the in-flight turn instead of the whole session. The | |
| 512 | + | // baseline counter advances in lock-step so the exit-time flush | |
| 513 | + | // doesn't double-write rows already on disk. | |
| 514 | + | persistMessage := func(m provider.Message) { | |
| 515 | + | persistMu.Lock() | |
| 516 | + | defer persistMu.Unlock() | |
| 517 | + | if sess == nil { | |
| 518 | + | return | |
| 519 | + | } | |
| 520 | + | if err := sess.AppendMessage(m); err == nil { | |
| 521 | + | sessBaselineMsgs++ | |
| 522 | + | } | |
| 523 | + | } | |
| 524 | + | persistUsage := func(cum provider.Usage) { | |
| 525 | + | persistMu.Lock() | |
| 526 | + | defer persistMu.Unlock() | |
| 527 | + | if sess == nil { | |
| 528 | + | return | |
| 529 | + | } | |
| 530 | + | _ = sess.AppendUsage(cum, cum) | |
| 531 | + | } | |
| 532 | + | wireAgentPersist := func(a *core.Agent) *core.Agent { | |
| 533 | + | if a == nil { | |
| 534 | + | return a | |
| 535 | + | } | |
| 536 | + | a.OnMessageAppended = persistMessage | |
| 537 | + | a.OnUsage = persistUsage | |
| 538 | + | return a | |
| 539 | + | } | |
| 540 | + | wireAgentPersist(ag) | |
| 541 | + | ||
| 542 | + | // Re-wrap the build closures so any agent constructed by the TUI | |
| 543 | + | // (login, /model swap to a different provider) also gets the | |
| 544 | + | // persistence hooks. Without this, switching provider would | |
| 545 | + | // silently revert to the old in-memory-only behaviour. | |
| 546 | + | baseBuildAgent := buildAgent | |
| 547 | + | buildAgent = func() (*core.Agent, string, string, error) { | |
| 548 | + | a, p, m, err := baseBuildAgent() | |
| 549 | + | return wireAgentPersist(a), p, m, err | |
| 550 | + | } | |
| 551 | + | baseBuildAgentFor := buildAgentFor | |
| 552 | + | buildAgentFor = func(providerOverride, modelOverride string) (*core.Agent, string, string, error) { | |
| 553 | + | a, p, m, err := baseBuildAgentFor(providerOverride, modelOverride) | |
| 554 | + | return wireAgentPersist(a), p, m, err | |
| 555 | + | } | |
| 556 | + | ||
| 497 | 557 | // loadSession replaces the current session with the one at path and | |
| 498 | 558 | // hands its messages to the agent. Used by the /sessions picker. | |
| 499 | 559 | loadSession := func(path string) error { | |
| @@ -505,9 +565,14 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 505 | 565 | if err != nil { | |
| 506 | 566 | return err | |
| 507 | 567 | } | |
| 568 | + | persistMu.Lock() | |
| 508 | 569 | // Flush any unsaved messages to the old session before swapping. | |
| 570 | + | // Per-message persistence keeps sessBaselineMsgs current, so | |
| 571 | + | // this is a defensive no-op in the common case; it still | |
| 572 | + | // matters for the rare race where a turn just finished and | |
| 573 | + | // hadn't fired its hook yet. | |
| 509 | 574 | if sess != nil { | |
| 510 | - | WriteNewTranscript(currentAg, sess, sessBaselineMsgs) | |
| 575 | + | writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs) | |
| 511 | 576 | _ = sess.Close() | |
| 512 | 577 | } | |
| 513 | 578 | sess = newSess | |
| @@ -516,6 +581,7 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 516 | 581 | currentAg.SeedCost(usage) | |
| 517 | 582 | } | |
| 518 | 583 | sessBaselineMsgs = len(msgs) | |
| 584 | + | persistMu.Unlock() | |
| 519 | 585 | return nil | |
| 520 | 586 | } | |
| 521 | 587 | ||
| @@ -612,18 +678,21 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 612 | 678 | // Append any not-yet-persisted agent messages to the | |
| 613 | 679 | // current session file, then advance the baseline so | |
| 614 | 680 | // the final WriteNewTranscript at exit doesn't write | |
| 615 | - | // duplicates. Called by /session export so the | |
| 616 | - | // exported bytes include everything the user has | |
| 617 | - | // said in the running TUI, not just whatever got | |
| 618 | - | // flushed at startup. | |
| 619 | - | if sess == nil { | |
| 620 | - | return | |
| 621 | - | } | |
| 681 | + | // duplicates. Per-message persistence keeps the on- | |
| 682 | + | // disk file current already, so this is mostly a | |
| 683 | + | // defensive flush — still needed for /session export | |
| 684 | + | // to guarantee the exported bytes include the very | |
| 685 | + | // last in-flight turn. | |
| 622 | 686 | currentAg := iv.Agent() | |
| 623 | 687 | if currentAg == nil { | |
| 624 | 688 | return | |
| 625 | 689 | } | |
| 626 | - | WriteNewTranscript(currentAg, sess, sessBaselineMsgs) | |
| 690 | + | persistMu.Lock() | |
| 691 | + | defer persistMu.Unlock() | |
| 692 | + | if sess == nil { | |
| 693 | + | return | |
| 694 | + | } | |
| 695 | + | writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs) | |
| 627 | 696 | sessBaselineMsgs = len(currentAg.Messages()) | |
| 628 | 697 | }, | |
| 629 | 698 | Extensions: extMgr, | |
| @@ -680,11 +749,52 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 680 | 749 | confirmGate.SetConfirmer(iv) | |
| 681 | 750 | } | |
| 682 | 751 | ||
| 752 | + | // Signal-driven flush: a SIGTERM / SIGHUP to the zot process | |
| 753 | + | // (closed terminal window, system shutdown, kill) used to lose | |
| 754 | + | // the entire in-memory transcript because the deferred post-Run | |
| 755 | + | // flush below never ran. Per-message persistence above covers | |
| 756 | + | // most of it; this handler writes any in-flight remainder and | |
| 757 | + | // then exits the process so we don't double-paint over a | |
| 758 | + | // broken terminal that the TUI's restore deferreds can no | |
| 759 | + | // longer fix from a signal context. | |
| 760 | + | // | |
| 761 | + | // SIGINT is intentionally NOT handled here — the TUI consumes | |
| 762 | + | // Ctrl+C as a regular key event for cancel/clear semantics, and | |
| 763 | + | // installing a SIGINT notifier here would swallow it. | |
| 764 | + | sigCh := make(chan os.Signal, 1) | |
| 765 | + | signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGHUP) | |
| 766 | + | defer signal.Stop(sigCh) | |
| 767 | + | go func() { | |
| 768 | + | _, ok := <-sigCh | |
| 769 | + | if !ok { | |
| 770 | + | return | |
| 771 | + | } | |
| 772 | + | if finalAg := iv.Agent(); finalAg != nil { | |
| 773 | + | persistMu.Lock() | |
| 774 | + | if sess != nil { | |
| 775 | + | writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs) | |
| 776 | + | sessBaselineMsgs = len(finalAg.Messages()) | |
| 777 | + | _ = sess.Close() | |
| 778 | + | sess = nil | |
| 779 | + | } | |
| 780 | + | persistMu.Unlock() | |
| 781 | + | } | |
| 782 | + | // Exit cleanly. Re-raising the signal would skip os.Exit's | |
| 783 | + | // at-exit hooks; explicit exit is fine because we've already | |
| 784 | + | // flushed the only at-risk state (the session file). | |
| 785 | + | os.Exit(0) | |
| 786 | + | }() | |
| 787 | + | ||
| 683 | 788 | runErr := iv.Run(ctx) | |
| 684 | 789 | ||
| 685 | 790 | // Flush final transcript to session (only if we had / ended up with an agent). | |
| 686 | - | if finalAg := iv.Agent(); finalAg != nil && sess != nil { | |
| 687 | - | WriteNewTranscript(finalAg, sess, sessBaselineMsgs) | |
| 791 | + | if finalAg := iv.Agent(); finalAg != nil { | |
| 792 | + | persistMu.Lock() | |
| 793 | + | if sess != nil { | |
| 794 | + | writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs) | |
| 795 | + | sessBaselineMsgs = len(finalAg.Messages()) | |
| 796 | + | } | |
| 797 | + | persistMu.Unlock() | |
| 688 | 798 | } | |
| 689 | 799 | return runErr | |
| 690 | 800 | } | |
| @@ -755,8 +865,18 @@ func pickSession(cwd string) (string, error) { | |||
| 755 | 865 | } | |
| 756 | 866 | ||
| 757 | 867 | // WriteNewTranscript appends only messages after index `from` from the | |
| 758 | - | // agent's transcript to the session. | |
| 868 | + | // agent's transcript to the session. Used by callers that don't hold | |
| 869 | + | // the persistMu (non-interactive print/json modes which run a single | |
| 870 | + | // turn under their own goroutine). | |
| 759 | 871 | func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) { | |
| 872 | + | writeNewTranscriptLocked(ag, sess, from) | |
| 873 | + | } | |
| 874 | + | ||
| 875 | + | // writeNewTranscriptLocked is the same as WriteNewTranscript. The | |
| 876 | + | // suffix marks that interactive callers must hold persistMu when | |
| 877 | + | // invoking it so concurrent appends from the agent loop don't race | |
| 878 | + | // with this catch-up flush. | |
| 879 | + | func writeNewTranscriptLocked(ag *core.Agent, sess *core.Session, from int) { | |
| 760 | 880 | if sess == nil || ag == nil { | |
| 761 | 881 | return | |
| 762 | 882 | } | |
| @@ -7,7 +7,10 @@ import ( | |||||
| 7 | "fmt" | 7 | "fmt" | ||
| 8 | "io" | 8 | "io" | ||
| 9 | "os" | 9 | "os" | ||
| 10 | + | "os/signal" | |||
| 10 | "strings" | 11 | "strings" | ||
| 12 | + | "sync" | |||
| 13 | + | "syscall" | |||
| 11 | "time" | 14 | "time" | ||
| 12 | 15 | ||||
| 13 | "github.com/patriceckhart/zot/internal/agent/extensions" | 16 | "github.com/patriceckhart/zot/internal/agent/extensions" | ||
| @@ -482,6 +485,12 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 482 | 485 | ||||
| 483 | var sess *core.Session | 486 | var sess *core.Session | ||
| 484 | var sessBaselineMsgs int // messages already on disk when current session opened | 487 | var sessBaselineMsgs int // messages already on disk when current session opened | ||
| 488 | + | // persistMu guards sess + sessBaselineMsgs against concurrent access | |||
| 489 | + | // from the agent loop's per-message persistence hook (runs on the | |||
| 490 | + | // agent goroutine) and the TUI's session swap / flush callbacks | |||
| 491 | + | // (run on the TUI goroutine). Without this, a /sessions swap that | |||
| 492 | + | // races with a finishing turn could double-write or lose messages. | |||
| 493 | + | var persistMu sync.Mutex | |||
| 485 | if !args.NoSess && ag != nil { | 494 | if !args.NoSess && ag != nil { | ||
| 486 | sess, _ = openOrCreateSession(args, r, ag, version) | 495 | sess, _ = openOrCreateSession(args, r, ag, version) | ||
| 487 | if ag != nil { | 496 | if ag != nil { | ||
| @@ -489,11 +498,62 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 489 | } | 498 | } | ||
| 490 | } | 499 | } | ||
| 491 | defer func() { | 500 | defer func() { | ||
| 501 | + | persistMu.Lock() | |||
| 502 | + | defer persistMu.Unlock() | |||
| 492 | if sess != nil { | 503 | if sess != nil { | ||
| 493 | sess.Close() | 504 | sess.Close() | ||
| 494 | } | 505 | } | ||
| 495 | }() | 506 | }() | ||
| 496 | 507 | ||||
| 508 | + | // persistMessage is the per-message hook bound to the agent. It | |||
| 509 | + | // appends each new transcript message to the live session as soon | |||
| 510 | + | // as it lands, so a kill / closed terminal / OS crash costs at | |||
| 511 | + | // most the in-flight turn instead of the whole session. The | |||
| 512 | + | // baseline counter advances in lock-step so the exit-time flush | |||
| 513 | + | // doesn't double-write rows already on disk. | |||
| 514 | + | persistMessage := func(m provider.Message) { | |||
| 515 | + | persistMu.Lock() | |||
| 516 | + | defer persistMu.Unlock() | |||
| 517 | + | if sess == nil { | |||
| 518 | + | return | |||
| 519 | + | } | |||
| 520 | + | if err := sess.AppendMessage(m); err == nil { | |||
| 521 | + | sessBaselineMsgs++ | |||
| 522 | + | } | |||
| 523 | + | } | |||
| 524 | + | persistUsage := func(cum provider.Usage) { | |||
| 525 | + | persistMu.Lock() | |||
| 526 | + | defer persistMu.Unlock() | |||
| 527 | + | if sess == nil { | |||
| 528 | + | return | |||
| 529 | + | } | |||
| 530 | + | _ = sess.AppendUsage(cum, cum) | |||
| 531 | + | } | |||
| 532 | + | wireAgentPersist := func(a *core.Agent) *core.Agent { | |||
| 533 | + | if a == nil { | |||
| 534 | + | return a | |||
| 535 | + | } | |||
| 536 | + | a.OnMessageAppended = persistMessage | |||
| 537 | + | a.OnUsage = persistUsage | |||
| 538 | + | return a | |||
| 539 | + | } | |||
| 540 | + | wireAgentPersist(ag) | |||
| 541 | + | ||||
| 542 | + | // Re-wrap the build closures so any agent constructed by the TUI | |||
| 543 | + | // (login, /model swap to a different provider) also gets the | |||
| 544 | + | // persistence hooks. Without this, switching provider would | |||
| 545 | + | // silently revert to the old in-memory-only behaviour. | |||
| 546 | + | baseBuildAgent := buildAgent | |||
| 547 | + | buildAgent = func() (*core.Agent, string, string, error) { | |||
| 548 | + | a, p, m, err := baseBuildAgent() | |||
| 549 | + | return wireAgentPersist(a), p, m, err | |||
| 550 | + | } | |||
| 551 | + | baseBuildAgentFor := buildAgentFor | |||
| 552 | + | buildAgentFor = func(providerOverride, modelOverride string) (*core.Agent, string, string, error) { | |||
| 553 | + | a, p, m, err := baseBuildAgentFor(providerOverride, modelOverride) | |||
| 554 | + | return wireAgentPersist(a), p, m, err | |||
| 555 | + | } | |||
| 556 | + | ||||
| 497 | // loadSession replaces the current session with the one at path and | 557 | // loadSession replaces the current session with the one at path and | ||
| 498 | // hands its messages to the agent. Used by the /sessions picker. | 558 | // hands its messages to the agent. Used by the /sessions picker. | ||
| 499 | loadSession := func(path string) error { | 559 | loadSession := func(path string) error { | ||
| @@ -505,9 +565,14 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 505 | if err != nil { | 565 | if err != nil { | ||
| 506 | return err | 566 | return err | ||
| 507 | } | 567 | } | ||
| 568 | + | persistMu.Lock() | |||
| 508 | // Flush any unsaved messages to the old session before swapping. | 569 | // Flush any unsaved messages to the old session before swapping. | ||
| 570 | + | // Per-message persistence keeps sessBaselineMsgs current, so | |||
| 571 | + | // this is a defensive no-op in the common case; it still | |||
| 572 | + | // matters for the rare race where a turn just finished and | |||
| 573 | + | // hadn't fired its hook yet. | |||
| 509 | if sess != nil { | 574 | if sess != nil { | ||
| 510 | - | WriteNewTranscript(currentAg, sess, sessBaselineMsgs) | 575 | + | writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs) |
| 511 | _ = sess.Close() | 576 | _ = sess.Close() | ||
| 512 | } | 577 | } | ||
| 513 | sess = newSess | 578 | sess = newSess | ||
| @@ -516,6 +581,7 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 516 | currentAg.SeedCost(usage) | 581 | currentAg.SeedCost(usage) | ||
| 517 | } | 582 | } | ||
| 518 | sessBaselineMsgs = len(msgs) | 583 | sessBaselineMsgs = len(msgs) | ||
| 584 | + | persistMu.Unlock() | |||
| 519 | return nil | 585 | return nil | ||
| 520 | } | 586 | } | ||
| 521 | 587 | ||||
| @@ -612,18 +678,21 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 612 | // Append any not-yet-persisted agent messages to the | 678 | // Append any not-yet-persisted agent messages to the | ||
| 613 | // current session file, then advance the baseline so | 679 | // current session file, then advance the baseline so | ||
| 614 | // the final WriteNewTranscript at exit doesn't write | 680 | // the final WriteNewTranscript at exit doesn't write | ||
| 615 | - | // duplicates. Called by /session export so the | 681 | + | // duplicates. Per-message persistence keeps the on- |
| 616 | - | // exported bytes include everything the user has | 682 | + | // disk file current already, so this is mostly a |
| 617 | - | // said in the running TUI, not just whatever got | 683 | + | // defensive flush — still needed for /session export |
| 618 | - | // flushed at startup. | 684 | + | // to guarantee the exported bytes include the very |
| 619 | - | if sess == nil { | 685 | + | // last in-flight turn. |
| 620 | - | return | |||
| 621 | - | } | |||
| 622 | currentAg := iv.Agent() | 686 | currentAg := iv.Agent() | ||
| 623 | if currentAg == nil { | 687 | if currentAg == nil { | ||
| 624 | return | 688 | return | ||
| 625 | } | 689 | } | ||
| 626 | - | WriteNewTranscript(currentAg, sess, sessBaselineMsgs) | 690 | + | persistMu.Lock() |
| 691 | + | defer persistMu.Unlock() | |||
| 692 | + | if sess == nil { | |||
| 693 | + | return | |||
| 694 | + | } | |||
| 695 | + | writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs) | |||
| 627 | sessBaselineMsgs = len(currentAg.Messages()) | 696 | sessBaselineMsgs = len(currentAg.Messages()) | ||
| 628 | }, | 697 | }, | ||
| 629 | Extensions: extMgr, | 698 | Extensions: extMgr, | ||
| @@ -680,11 +749,52 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 680 | confirmGate.SetConfirmer(iv) | 749 | confirmGate.SetConfirmer(iv) | ||
| 681 | } | 750 | } | ||
| 682 | 751 | ||||
| 752 | + | // Signal-driven flush: a SIGTERM / SIGHUP to the zot process | |||
| 753 | + | // (closed terminal window, system shutdown, kill) used to lose | |||
| 754 | + | // the entire in-memory transcript because the deferred post-Run | |||
| 755 | + | // flush below never ran. Per-message persistence above covers | |||
| 756 | + | // most of it; this handler writes any in-flight remainder and | |||
| 757 | + | // then exits the process so we don't double-paint over a | |||
| 758 | + | // broken terminal that the TUI's restore deferreds can no | |||
| 759 | + | // longer fix from a signal context. | |||
| 760 | + | // | |||
| 761 | + | // SIGINT is intentionally NOT handled here — the TUI consumes | |||
| 762 | + | // Ctrl+C as a regular key event for cancel/clear semantics, and | |||
| 763 | + | // installing a SIGINT notifier here would swallow it. | |||
| 764 | + | sigCh := make(chan os.Signal, 1) | |||
| 765 | + | signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGHUP) | |||
| 766 | + | defer signal.Stop(sigCh) | |||
| 767 | + | go func() { | |||
| 768 | + | _, ok := <-sigCh | |||
| 769 | + | if !ok { | |||
| 770 | + | return | |||
| 771 | + | } | |||
| 772 | + | if finalAg := iv.Agent(); finalAg != nil { | |||
| 773 | + | persistMu.Lock() | |||
| 774 | + | if sess != nil { | |||
| 775 | + | writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs) | |||
| 776 | + | sessBaselineMsgs = len(finalAg.Messages()) | |||
| 777 | + | _ = sess.Close() | |||
| 778 | + | sess = nil | |||
| 779 | + | } | |||
| 780 | + | persistMu.Unlock() | |||
| 781 | + | } | |||
| 782 | + | // Exit cleanly. Re-raising the signal would skip os.Exit's | |||
| 783 | + | // at-exit hooks; explicit exit is fine because we've already | |||
| 784 | + | // flushed the only at-risk state (the session file). | |||
| 785 | + | os.Exit(0) | |||
| 786 | + | }() | |||
| 787 | + | ||||
| 683 | runErr := iv.Run(ctx) | 788 | runErr := iv.Run(ctx) | ||
| 684 | 789 | ||||
| 685 | // Flush final transcript to session (only if we had / ended up with an agent). | 790 | // Flush final transcript to session (only if we had / ended up with an agent). | ||
| 686 | - | if finalAg := iv.Agent(); finalAg != nil && sess != nil { | 791 | + | if finalAg := iv.Agent(); finalAg != nil { |
| 687 | - | WriteNewTranscript(finalAg, sess, sessBaselineMsgs) | 792 | + | persistMu.Lock() |
| 793 | + | if sess != nil { | |||
| 794 | + | writeNewTranscriptLocked(finalAg, sess, sessBaselineMsgs) | |||
| 795 | + | sessBaselineMsgs = len(finalAg.Messages()) | |||
| 796 | + | } | |||
| 797 | + | persistMu.Unlock() | |||
| 688 | } | 798 | } | ||
| 689 | return runErr | 799 | return runErr | ||
| 690 | } | 800 | } | ||
| @@ -755,8 +865,18 @@ func pickSession(cwd string) (string, error) { | |||||
| 755 | } | 865 | } | ||
| 756 | 866 | ||||
| 757 | // WriteNewTranscript appends only messages after index `from` from the | 867 | // WriteNewTranscript appends only messages after index `from` from the | ||
| 758 | - | // agent's transcript to the session. | 868 | + | // agent's transcript to the session. Used by callers that don't hold |
| 869 | + | // the persistMu (non-interactive print/json modes which run a single | |||
| 870 | + | // turn under their own goroutine). | |||
| 759 | func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) { | 871 | func WriteNewTranscript(ag *core.Agent, sess *core.Session, from int) { | ||
| 872 | + | writeNewTranscriptLocked(ag, sess, from) | |||
| 873 | + | } | |||
| 874 | + | ||||
| 875 | + | // writeNewTranscriptLocked is the same as WriteNewTranscript. The | |||
| 876 | + | // suffix marks that interactive callers must hold persistMu when | |||
| 877 | + | // invoking it so concurrent appends from the agent loop don't race | |||
| 878 | + | // with this catch-up flush. | |||
| 879 | + | func writeNewTranscriptLocked(ag *core.Agent, sess *core.Session, from int) { | |||
| 760 | if sess == nil || ag == nil { | 880 | if sess == nil || ag == nil { | ||
| 761 | return | 881 | return | ||
| 762 | } | 882 | } | ||
M internal/core/agent.go +36 -0
| @@ -51,6 +51,21 @@ type Agent struct { | |||
| 51 | 51 | // without each caller having to compose sinks manually. | |
| 52 | 52 | OnEvent func(AgentEvent) | |
| 53 | 53 | ||
| 54 | + | // OnMessageAppended, if set, fires every time a message is | |
| 55 | + | // appended to the in-memory transcript by the agent loop — the | |
| 56 | + | // initial user prompt, each finalised assistant message, and | |
| 57 | + | // each tool-results message (plus the synthetic OpenAI image | |
| 58 | + | // mirror, if any). Hosts wire this to the on-disk session so | |
| 59 | + | // that turns are durable as soon as they happen, instead of | |
| 60 | + | // only being flushed on a clean exit. | |
| 61 | + | OnMessageAppended func(provider.Message) | |
| 62 | + | ||
| 63 | + | // OnUsage, if set, fires after every turn's usage row arrives, | |
| 64 | + | // carrying the cumulative usage for the session. Hosts wire | |
| 65 | + | // this to the on-disk session so the persisted total stays | |
| 66 | + | // current and a crash recovers the right cost figure. | |
| 67 | + | OnUsage func(cumulative provider.Usage) | |
| 68 | + | ||
| 54 | 69 | mu sync.Mutex | |
| 55 | 70 | messages []provider.Message | |
| 56 | 71 | cost CostTracker | |
| @@ -108,6 +123,16 @@ func (a *Agent) SeedCost(u provider.Usage) { | |||
| 108 | 123 | a.cost.Total = u | |
| 109 | 124 | } | |
| 110 | 125 | ||
| 126 | + | // fireMessageAppended invokes OnMessageAppended without holding the | |
| 127 | + | // agent mutex, so the host's persistence callback can take its own | |
| 128 | + | // locks without deadlocking the agent loop. Tolerates a nil hook so | |
| 129 | + | // non-persisting callers (tests, RPC mode) don't have to set it. | |
| 130 | + | func (a *Agent) fireMessageAppended(m provider.Message) { | |
| 131 | + | if a.OnMessageAppended != nil { | |
| 132 | + | a.OnMessageAppended(m) | |
| 133 | + | } | |
| 134 | + | } | |
| 135 | + | ||
| 111 | 136 | // Prompt sends a user message and runs the agent loop until the model | |
| 112 | 137 | // stops or an error occurs. Events are delivered via sink in order. | |
| 113 | 138 | // sink must not block the caller for long; buffer as needed. | |
| @@ -128,6 +153,7 @@ func (a *Agent) Prompt(ctx context.Context, text string, images []provider.Image | |||
| 128 | 153 | a.mu.Lock() | |
| 129 | 154 | a.messages = append(a.messages, user) | |
| 130 | 155 | a.mu.Unlock() | |
| 156 | + | a.fireMessageAppended(user) | |
| 131 | 157 | sink(EvUserMessage{Message: user}) | |
| 132 | 158 | ||
| 133 | 159 | return a.runLoop(ctx, sink) | |
| @@ -188,12 +214,18 @@ func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error { | |||
| 188 | 214 | // This keeps the transcript self-contained for providers that can | |
| 189 | 215 | // see image blocks in tool messages while making OpenAI vision | |
| 190 | 216 | // models actually receive the image bytes. | |
| 217 | + | var imageMirror provider.Message | |
| 191 | 218 | if a.Client != nil && a.Client.Name() == "openai" { | |
| 192 | 219 | if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 { | |
| 193 | 220 | a.messages = append(a.messages, mirror) | |
| 221 | + | imageMirror = mirror | |
| 194 | 222 | } | |
| 195 | 223 | } | |
| 196 | 224 | a.mu.Unlock() | |
| 225 | + | a.fireMessageAppended(toolMsg) | |
| 226 | + | if len(imageMirror.Content) > 0 { | |
| 227 | + | a.fireMessageAppended(imageMirror) | |
| 228 | + | } | |
| 197 | 229 | // If context was cancelled during tool execution, bail out. | |
| 198 | 230 | if err := ctx.Err(); err != nil { | |
| 199 | 231 | sink(EvDone{}) | |
| @@ -252,6 +284,9 @@ func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.St | |||
| 252 | 284 | case provider.EventUsage: | |
| 253 | 285 | cum := a.cost.Add(e.Usage) | |
| 254 | 286 | sink(EvUsage{Usage: e.Usage, Cumulative: cum}) | |
| 287 | + | if a.OnUsage != nil { | |
| 288 | + | a.OnUsage(cum) | |
| 289 | + | } | |
| 255 | 290 | case provider.EventDone: | |
| 256 | 291 | stop = e.Stop | |
| 257 | 292 | finalErr = e.Err | |
| @@ -302,6 +337,7 @@ func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.St | |||
| 302 | 337 | a.mu.Lock() | |
| 303 | 338 | a.messages = append(a.messages, finalMsg) | |
| 304 | 339 | a.mu.Unlock() | |
| 340 | + | a.fireMessageAppended(finalMsg) | |
| 305 | 341 | if !suppress { | |
| 306 | 342 | sink(EvAssistantMessage{Message: emit}) | |
| 307 | 343 | } | |
| @@ -51,6 +51,21 @@ type Agent struct { | |||||
| 51 | // without each caller having to compose sinks manually. | 51 | // without each caller having to compose sinks manually. | ||
| 52 | OnEvent func(AgentEvent) | 52 | OnEvent func(AgentEvent) | ||
| 53 | 53 | ||||
| 54 | + | // OnMessageAppended, if set, fires every time a message is | |||
| 55 | + | // appended to the in-memory transcript by the agent loop — the | |||
| 56 | + | // initial user prompt, each finalised assistant message, and | |||
| 57 | + | // each tool-results message (plus the synthetic OpenAI image | |||
| 58 | + | // mirror, if any). Hosts wire this to the on-disk session so | |||
| 59 | + | // that turns are durable as soon as they happen, instead of | |||
| 60 | + | // only being flushed on a clean exit. | |||
| 61 | + | OnMessageAppended func(provider.Message) | |||
| 62 | + | ||||
| 63 | + | // OnUsage, if set, fires after every turn's usage row arrives, | |||
| 64 | + | // carrying the cumulative usage for the session. Hosts wire | |||
| 65 | + | // this to the on-disk session so the persisted total stays | |||
| 66 | + | // current and a crash recovers the right cost figure. | |||
| 67 | + | OnUsage func(cumulative provider.Usage) | |||
| 68 | + | ||||
| 54 | mu sync.Mutex | 69 | mu sync.Mutex | ||
| 55 | messages []provider.Message | 70 | messages []provider.Message | ||
| 56 | cost CostTracker | 71 | cost CostTracker | ||
| @@ -108,6 +123,16 @@ func (a *Agent) SeedCost(u provider.Usage) { | |||||
| 108 | a.cost.Total = u | 123 | a.cost.Total = u | ||
| 109 | } | 124 | } | ||
| 110 | 125 | ||||
| 126 | + | // fireMessageAppended invokes OnMessageAppended without holding the | |||
| 127 | + | // agent mutex, so the host's persistence callback can take its own | |||
| 128 | + | // locks without deadlocking the agent loop. Tolerates a nil hook so | |||
| 129 | + | // non-persisting callers (tests, RPC mode) don't have to set it. | |||
| 130 | + | func (a *Agent) fireMessageAppended(m provider.Message) { | |||
| 131 | + | if a.OnMessageAppended != nil { | |||
| 132 | + | a.OnMessageAppended(m) | |||
| 133 | + | } | |||
| 134 | + | } | |||
| 135 | + | ||||
| 111 | // Prompt sends a user message and runs the agent loop until the model | 136 | // Prompt sends a user message and runs the agent loop until the model | ||
| 112 | // stops or an error occurs. Events are delivered via sink in order. | 137 | // stops or an error occurs. Events are delivered via sink in order. | ||
| 113 | // sink must not block the caller for long; buffer as needed. | 138 | // sink must not block the caller for long; buffer as needed. | ||
| @@ -128,6 +153,7 @@ func (a *Agent) Prompt(ctx context.Context, text string, images []provider.Image | |||||
| 128 | a.mu.Lock() | 153 | a.mu.Lock() | ||
| 129 | a.messages = append(a.messages, user) | 154 | a.messages = append(a.messages, user) | ||
| 130 | a.mu.Unlock() | 155 | a.mu.Unlock() | ||
| 156 | + | a.fireMessageAppended(user) | |||
| 131 | sink(EvUserMessage{Message: user}) | 157 | sink(EvUserMessage{Message: user}) | ||
| 132 | 158 | ||||
| 133 | return a.runLoop(ctx, sink) | 159 | return a.runLoop(ctx, sink) | ||
| @@ -188,12 +214,18 @@ func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error { | |||||
| 188 | // This keeps the transcript self-contained for providers that can | 214 | // This keeps the transcript self-contained for providers that can | ||
| 189 | // see image blocks in tool messages while making OpenAI vision | 215 | // see image blocks in tool messages while making OpenAI vision | ||
| 190 | // models actually receive the image bytes. | 216 | // models actually receive the image bytes. | ||
| 217 | + | var imageMirror provider.Message | |||
| 191 | if a.Client != nil && a.Client.Name() == "openai" { | 218 | if a.Client != nil && a.Client.Name() == "openai" { | ||
| 192 | if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 { | 219 | if mirror := mirrorToolImagesAsUser(toolMsg); len(mirror.Content) > 0 { | ||
| 193 | a.messages = append(a.messages, mirror) | 220 | a.messages = append(a.messages, mirror) | ||
| 221 | + | imageMirror = mirror | |||
| 194 | } | 222 | } | ||
| 195 | } | 223 | } | ||
| 196 | a.mu.Unlock() | 224 | a.mu.Unlock() | ||
| 225 | + | a.fireMessageAppended(toolMsg) | |||
| 226 | + | if len(imageMirror.Content) > 0 { | |||
| 227 | + | a.fireMessageAppended(imageMirror) | |||
| 228 | + | } | |||
| 197 | // If context was cancelled during tool execution, bail out. | 229 | // If context was cancelled during tool execution, bail out. | ||
| 198 | if err := ctx.Err(); err != nil { | 230 | if err := ctx.Err(); err != nil { | ||
| 199 | sink(EvDone{}) | 231 | sink(EvDone{}) | ||
| @@ -252,6 +284,9 @@ func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.St | |||||
| 252 | case provider.EventUsage: | 284 | case provider.EventUsage: | ||
| 253 | cum := a.cost.Add(e.Usage) | 285 | cum := a.cost.Add(e.Usage) | ||
| 254 | sink(EvUsage{Usage: e.Usage, Cumulative: cum}) | 286 | sink(EvUsage{Usage: e.Usage, Cumulative: cum}) | ||
| 287 | + | if a.OnUsage != nil { | |||
| 288 | + | a.OnUsage(cum) | |||
| 289 | + | } | |||
| 255 | case provider.EventDone: | 290 | case provider.EventDone: | ||
| 256 | stop = e.Stop | 291 | stop = e.Stop | ||
| 257 | finalErr = e.Err | 292 | finalErr = e.Err | ||
| @@ -302,6 +337,7 @@ func (a *Agent) oneTurn(ctx context.Context, sink func(AgentEvent)) (provider.St | |||||
| 302 | a.mu.Lock() | 337 | a.mu.Lock() | ||
| 303 | a.messages = append(a.messages, finalMsg) | 338 | a.messages = append(a.messages, finalMsg) | ||
| 304 | a.mu.Unlock() | 339 | a.mu.Unlock() | ||
| 340 | + | a.fireMessageAppended(finalMsg) | |||
| 305 | if !suppress { | 341 | if !suppress { | ||
| 306 | sink(EvAssistantMessage{Message: emit}) | 342 | sink(EvAssistantMessage{Message: emit}) | ||
| 307 | } | 343 | } | ||