display fixes: markdown rendering, tool-specific components

1 files changed, 228 insertions, 17 deletions

diff --git a/src/tui_app.zig b/src/tui_app.zig
index f9c58fe..e44a74a 100644
--- a/src/tui_app.zig
+++ b/src/tui_app.zig
@@ -67,6 +67,7 @@ const selectors_mod = @import("tui_selectors.zig");
 const config_file = @import("config_file.zig");
 const auth_manager = @import("auth_manager.zig");
 const models_toml = @import("models_toml.zig");
+const pricing_format = @import("pricing_format.zig");
 const tui_key = @import("tui_key.zig");
 
 const Terminal = terminal_mod.Terminal;
@@ -301,6 +302,15 @@ pub const App = struct {
     /// the picker overlays.
     selectors: ?*SelectorController = null,
 
+    /// Optional hook the App invokes on every `message_complete` carrying
+    /// a `Usage`. The hook (installed by the `SelectorController`) updates
+    /// session-running totals keyed by the current `(provider, model)` and
+    /// pushes the latest values into the footer. With no hook installed
+    /// (e.g. tests) the App still updates the per-turn context-window
+    /// tokens but does NOT accumulate session totals.
+    usage_record_ctx: ?*anyopaque = null,
+    usage_record_fn: ?*const fn (ctx: *anyopaque, usage: panto.Usage) void = null,
+
     /// Whether the input box currently participates in the engine list. It is
     /// removed during an in-flight turn (so streaming output appends below the
     /// transcript) and re-added when the turn completes. P1 keeps it simple:
@@ -358,6 +368,21 @@ pub const App = struct {
         self.flush_fn = f;
     }
 
+    /// Install the per-turn usage record hook (the `SelectorController`
+    /// calls this). On every `message_complete` the App invokes the
+    /// hook with the just-reported `Usage`; the hook keys the usage by
+    /// the current `(provider, model)` (its own concern) and pushes the
+    /// new session totals back into the footer. Tests omit this; the
+    /// per-turn context-window tokens still get updated.
+    pub fn setUsageRecorder(
+        self: *App,
+        ctx: *anyopaque,
+        f: *const fn (ctx: *anyopaque, usage: panto.Usage) void,
+    ) void {
+        self.usage_record_ctx = ctx;
+        self.usage_record_fn = f;
+    }
+
     fn flushSink(self: *App) void {
         if (self.flush_fn) |f| f(self.flush_ctx.?);
     }
@@ -866,6 +891,10 @@ pub const App = struct {
             .block_complete => |b| {
                 switch (b.block) {
                     .Text => {
+                        if (self.router.get(b.index)) |ref| switch (ref) {
+                            .assistant => |box| box.finishStream(),
+                            else => {},
+                        };
                         try self.fireTextLifecycle(b.index, .assistant, .assistant_text_complete, "assistant_text_complete", .{ .assistant_text = .{
                             .index = b.index,
                             .text = if (self.router.get(b.index)) |r| (if (r == .assistant) r.assistant.buffer.items else "") else "",
@@ -915,8 +944,13 @@ pub const App = struct {
                 // (output/reasoning excluded — not "in the window"). Latest
                 // value wins; not accumulated.
                 if (mc.usage) |u| {
-                    const ctx = u.input + u.cache_read + u.cache_write + u.output;
+                    const ctx = u.input + u.cache_read + u.cache_write;
                     self.footer.setContextTokens(ctx);
+                    // Hand the raw usage to the recorder (the
+                    // `SelectorController` is the canonical recorder; it
+                    // knows which `(provider, model)` produced this turn
+                    // and accumulates per-model token + cost totals).
+                    if (self.usage_record_fn) |f| f(self.usage_record_ctx.?, u);
                     self.scheduler.requestRender();
                 }
             },
@@ -939,6 +973,11 @@ pub const App = struct {
                 }
                 self.scheduler.requestRender();
             },
+            .tool_dispatch_result => |info| {
+                // Eager per-tool result carrier. Correlate by tool_use_id just
+                // like the aggregate completion event.
+                try self.routeToolResults(info.message);
+            },
             .tool_dispatch_complete => |info| {
                 // ToolResult blocks are delivered together here as the content
                 // of the appended user message. Correlate each back to its
@@ -1019,6 +1058,20 @@ fn modelPickerRowLessThan(_: void, a: ModelPickerRow, b: ModelPickerRow) bool {
     return std.mem.lessThan(u8, a.item.detail, b.item.detail);
 }
 
+fn addTokenBucket(
+    alloc: std.mem.Allocator,
+    buckets: *std.StringHashMapUnmanaged(u64),
+    key: []const u8,
+    amount: u64,
+) void {
+    const gop = buckets.getOrPut(alloc, key) catch return;
+    if (!gop.found_existing) {
+        gop.key_ptr.* = alloc.dupe(u8, key) catch return;
+        gop.value_ptr.* = 0;
+    }
+    gop.value_ptr.* +%= amount;
+}
+
 pub const SelectorController = struct {
     alloc: std.mem.Allocator,
     app: *App,
@@ -1027,6 +1080,11 @@ pub const SelectorController = struct {
     /// transport/auth and per-alias knobs lookups via `buildProviderConfig`.
     file_cfg: *const config_file.Config,
     defs: *const models_toml.ModelRegistry,
+    /// Per-(provider, wire-model) pricing table for the session. Borrowed
+    /// (the merged `models_toml.Models` outlives the controller). The
+    /// model-picker detail line and the footer session cost both look up
+    /// here; the controller does NOT mutate it.
+    pricing: *const panto.PricingRegistry,
     /// The live agent config snapshot, owned here. `agent` holds a pointer to
     /// it; we mutate `provider` in place and re-`setConfig` so the change is
     /// observed at the next turn.
@@ -1056,12 +1114,35 @@ pub const SelectorController = struct {
     /// The current model label ("provider:alias"), for the footer/preselect.
     model_label: []u8,
 
+    /// Session-running token totals, keyed by the `(provider, model)`
+    /// pair that produced them. A model switch in the middle of a
+    /// session just appends a new bucket; the SUM is what the footer
+    /// displays. Each entry is owned here and is the integer sum of
+    /// `input + output + cache_read + cache_write` for the turns
+    /// produced by that model.
+    ///
+    /// WHY PER-MODEL: the user can switch models mid-session; a
+    /// flat u64 would conflate two models' tokens. Per-model lets us
+    /// re-pricing: if the user pastes a corrected `models.toml` mid-
+    /// session we just rebuild the cost (the next addCost call picks
+    /// up the new pricing), and the token sum is the *same* flat
+    /// total regardless of pricing changes.
+    session_token_buckets: std.StringHashMapUnmanaged(u64) = .empty,
+
+    /// Session-running cost total in micro-cents. `null` means "at
+    /// least one priced component of at least one turn was unknown"
+    /// (the `addCost` poison rule). Set to 0 for the first
+    /// ALL-ZERO turn and stays known thereafter; any later
+    /// `costMicroCents == null` re-poisons.
+    session_cost: ?u64 = 0,
+
     pub fn init(
         alloc: std.mem.Allocator,
         app: *App,
         agent: *panto.Agent,
         file_cfg: *const config_file.Config,
         defs: *const models_toml.ModelRegistry,
+        pricing: *const panto.PricingRegistry,
         live: *panto.Config,
         initial_label: []const u8,
     ) !*SelectorController {
@@ -1073,12 +1154,17 @@ pub const SelectorController = struct {
             .agent = agent,
             .file_cfg = file_cfg,
             .defs = defs,
+            .pricing = pricing,
             .live = live,
             .model_items = &.{},
             .model_label = try alloc.dupe(u8, initial_label),
         };
         try self.buildModelItems();
         try self.refreshFooter();
+        // Install ourselves as the App's per-turn usage recorder so every
+        // `message_complete` lands in our session-running buckets and the
+        // footer session cost/token total is updated.
+        app.setUsageRecorder(self, recordUsageThunk);
         return self;
     }
 
@@ -1105,6 +1191,10 @@ pub const SelectorController = struct {
         self.alloc.free(self.model_entry_indices);
         self.alloc.free(self.reasoning_items);
         self.alloc.free(self.model_label);
+        // The session token buckets own the (provider, model) key strings.
+        var it = self.session_token_buckets.iterator();
+        while (it.next()) |entry| self.alloc.free(entry.key_ptr.*);
+        self.session_token_buckets.deinit(self.alloc);
         self.alloc.destroy(self);
     }
 
@@ -1116,7 +1206,7 @@ pub const SelectorController = struct {
         for (self.defs.entries.items, 0..) |d, def_index| {
             const label = try std.fmt.allocPrint(a, "{s}:{s}", .{ d.provider, d.alias });
             try self.model_strings.append(a, label);
-            const detail = try formatModelDetail(a, d);
+            const detail = try formatModelDetail(a, d, self.pricing);
             try self.model_strings.append(a, detail);
             try rows.append(a, .{ .def_index = def_index, .item = .{ .label = label, .detail = detail } });
         }
@@ -1246,14 +1336,91 @@ pub const SelectorController = struct {
         defer self.alloc.free(msg);
         _ = self.app.spawnStatus(msg) catch {};
     }
+
+    /// Record one turn's `usage` against the current `(provider, model)`,
+    /// then push the new session totals into the footer. Model-switch
+    /// tolerance: each `(provider, model)` is a separate bucket for
+    /// the per-model token sum; the FOOTER displays the flat sum
+    /// across all buckets, so a switch in the middle of a session is
+    /// invisible to the user (it just makes the next turn's tokens
+    /// land in a new bucket). For cost, the same `addCost` accumulator
+    /// is used; the per-turn cost is computed against THIS turn's
+    /// `(provider, model)` pricing, so a model with known pricing
+    /// before an unpriced model after still produces a known total
+    /// only up to the switch.
+    fn recordUsage(self: *SelectorController, usage: panto.Usage) void {
+        // Resolve the (provider, wire-model) pair for THIS turn. The
+        // pricing registry is keyed on these strings; the bucket
+        // string is "<provider>:<wire>" so a model name that happens
+        // to be reused across providers doesn't collide.
+        const colon = std.mem.indexOfScalar(u8, self.model_label, ':') orelse {
+            // A malformed label (no colon) is a programmer error in
+            // the boot sequence; just bail. The App still shows the
+            // per-turn context-window tokens.
+            return;
+        };
+        const provider_name = self.model_label[0..colon];
+        const wire_model = selectors_mod.wireModel(self.live.provider);
+
+        const turn_tokens = usage.input + usage.output + usage.cache_read + usage.cache_write;
+        const key = std.fmt.allocPrint(self.alloc, "{s}:{s}", .{ provider_name, wire_model }) catch return;
+        defer self.alloc.free(key);
+        addTokenBucket(self.alloc, &self.session_token_buckets, key, turn_tokens);
+
+        // Recompute the flat session sum. Cheap: the bucket count is
+        // tiny (the user doesn't switch models a thousand times).
+        var total: u64 = 0;
+        var it = self.session_token_buckets.iterator();
+        while (it.next()) |entry| total +%= entry.value_ptr.*;
+        self.app.footer.setSessionTokens(total);
+
+        // Cost: look up the pricing for the model that JUST produced
+        // this usage. If any priced component is null (no models.toml
+        // entry for this model), this turn's cost is null and
+        // `addCost` poisons the session total to null — once poisoned,
+        // it stays null for the rest of the session (no way to
+        // "un-poison" a `?u64` back to a known value).
+        //
+        // An empty (default-constructed) Pricing has every field null;
+        // `costMicroCents` treats every nonzero token usage as
+        // "unknown" and returns null. That's the same as the
+        // no-pricing case, so a single code path covers both.
+        const turn_cost: ?u64 = if (self.pricing.get(provider_name, wire_model)) |pricing|
+            panto.costMicroCents(usage, pricing)
+        else
+            null;
+        self.session_cost = panto.addCost(self.session_cost, turn_cost);
+        self.app.footer.setSessionCost(self.session_cost);
+    }
+
+    /// Static thunk for the App's `setUsageRecorder` callback. The
+    /// `ctx` we install with is the `*SelectorController` itself; the
+    /// thunk casts it back and dispatches to the typed method.
+    fn recordUsageThunk(ctx: *anyopaque, usage: panto.Usage) void {
+        const self: *SelectorController = @ptrCast(@alignCast(ctx));
+        self.recordUsage(usage);
+    }
 };
 
-/// Format the dim detail string for a model item: the wire model id plus the
-/// reasoning/thinking knobs declared in `models.toml`.
-fn formatModelDetail(alloc: std.mem.Allocator, d: models_toml.ModelDef) ![]u8 {
-    // Anthropic-style entries advertise thinking/effort; openai-style ones
-    // advertise reasoning. We don't know the provider's API style here, so we
-    // show whatever knobs are non-default.
+
+/// Format the dim detail string for a model item: the wire model id,
+/// the reasoning/thinking knobs declared in `models.toml`, and the
+/// pricing (when known). The pricing is looked up by the wire model
+/// id against the parsed `PricingRegistry` (the same keying the
+/// session cost accumulator uses); it surfaces as a compact
+/// `"1i/5o/0.1r/1.25w"` suffix on the same line so the picker row
+/// is a one-glance summary.
+///
+/// Anthropic-style entries advertise thinking/effort; openai-style
+/// ones advertise reasoning. We don't know the provider's API style
+/// here, so we show whatever knobs are non-default. Pricing is
+/// appended last so the most-novel information lands closest to the
+/// model label and the older knob info reads as supporting context.
+fn formatModelDetail(
+    alloc: std.mem.Allocator,
+    d: models_toml.ModelDef,
+    pricing: *const panto.PricingRegistry,
+) ![]u8 {
     var buf: std.ArrayList(u8) = .empty;
     errdefer buf.deinit(alloc);
     try buf.appendSlice(alloc, d.model);
@@ -1268,6 +1435,13 @@ fn formatModelDetail(alloc: std.mem.Allocator, d: models_toml.ModelDef) ![]u8 {
         try buf.appendSlice(alloc, "  reasoning:");
         try buf.appendSlice(alloc, @tagName(d.reasoning));
     }
+    if (pricing.get(d.provider, d.model)) |p| {
+        var scratch: [64]u8 = undefined;
+        if (pricing_format.formatPriceCompact(p, &scratch)) |tag| {
+            try buf.appendSlice(alloc, "  ");
+            try buf.appendSlice(alloc, tag);
+        }
+    }
     return buf.toOwnedSlice(alloc);
 }
 
@@ -2030,13 +2204,15 @@ test "routeEvent: full event stream renders through the real engine, no stdout"
     h.app.beginTurn();
     try h.app.routeEvent(.{ .message_start = .assistant });
     try h.app.routeEvent(.{ .block_start = .{ .block_type = .Text, .index = 0 } });
+    // The streaming path renders complete lines through markdown and shows the
+    // trailing partial line verbatim as it arrives.
     try h.app.routeEvent(delta(0, "Hi there"));
-    try h.app.routeEvent(.{ .turn_complete = {} });
-
     try h.app.renderNow();
-    const out = h.buf.written();
-    // The assistant text reached the engine output (not stdout).
-    try testing.expect(std.mem.indexOf(u8, out, "Hi there") != null);
+    try testing.expect(std.mem.indexOf(u8, h.buf.written(), "Hi there") != null);
+    try h.app.routeEvent(delta(0, "\n"));
+    try h.app.renderNow();
+    try testing.expect(std.mem.indexOf(u8, h.buf.written(), "Hi there") != null);
+    try h.app.routeEvent(.{ .turn_complete = {} });
 }
 
 test "beginTurn clears the block-index map but keeps transcript history" {
@@ -2306,12 +2482,17 @@ fn testModelDef(provider: []const u8, alias: []const u8, model: []const u8) mode
     };
 }
 
-test "formatModelDetail: shows wire model + non-default knobs" {
+test "formatModelDetail: shows wire model + non-default knobs + pricing" {
     const alloc = testing.allocator;
+    // Empty pricing registry: no tag appended (this is the default shape for
+    // test definitions that don't bother declaring a price).
+    var pricing = panto.PricingRegistry.init(alloc);
+    defer pricing.deinit();
+
     // Plain entry: just the wire model id.
     {
         const d = testModelDef("openai", "gpt", "gpt-4o");
-        const s = try formatModelDetail(alloc, d);
+        const s = try formatModelDetail(alloc, d, &pricing);
         defer alloc.free(s);
         try testing.expectEqualStrings("gpt-4o", s);
     }
@@ -2319,7 +2500,7 @@ test "formatModelDetail: shows wire model + non-default knobs" {
     {
         var d = testModelDef("openai", "o3", "o3");
         d.reasoning = .high;
-        const s = try formatModelDetail(alloc, d);
+        const s = try formatModelDetail(alloc, d, &pricing);
         defer alloc.free(s);
         try testing.expect(std.mem.indexOf(u8, s, "reasoning:high") != null);
     }
@@ -2328,11 +2509,24 @@ test "formatModelDetail: shows wire model + non-default knobs" {
         var d = testModelDef("anthropic", "opus", "claude-opus-4");
         d.thinking = .adaptive;
         d.effort = .xhigh;
-        const s = try formatModelDetail(alloc, d);
+        const s = try formatModelDetail(alloc, d, &pricing);
         defer alloc.free(s);
         try testing.expect(std.mem.indexOf(u8, s, "thinking:adaptive") != null);
         try testing.expect(std.mem.indexOf(u8, s, "effort:xhigh") != null);
     }
+    // With pricing: a tag like "1i/5o/0.1r/1.25w" is appended to the line.
+    {
+        try pricing.set("openai", "gpt-4o", .{
+            .input = 250,
+            .output = 1000,
+            .cache_read = 125,
+            .cache_write = 0,
+        });
+        const d = testModelDef("openai", "gpt-4o", "gpt-4o");
+        const s = try formatModelDetail(alloc, d, &pricing);
+        defer alloc.free(s);
+        try testing.expectEqualStrings("gpt-4o  2.5i/10o/1.25r/0w", s);
+    }
 }
 
 test "model picker rows sort by provider:alias and keep mapping" {
@@ -3083,3 +3277,20 @@ test "splitEditorArgv: splits flags, appends the path, and falls back to vi" {
         try testing.expectEqualStrings("/tmp/y.md", argv.items[1]);
     }
 }
+
+test "session token bucket helper initializes new buckets at zero" {
+    var buckets: std.StringHashMapUnmanaged(u64) = .empty;
+    defer {
+        var kit = buckets.keyIterator();
+        while (kit.next()) |k| testing.allocator.free(k.*);
+        buckets.deinit(testing.allocator);
+    }
+
+    addTokenBucket(testing.allocator, &buckets, "anthropic:haiku", 5);
+    addTokenBucket(testing.allocator, &buckets, "anthropic:haiku", 7);
+
+    var it = buckets.iterator();
+    const entry = it.next() orelse unreachable;
+    try testing.expectEqualStrings("anthropic:haiku", entry.key_ptr.*);
+    try testing.expectEqual(@as(u64, 12), entry.value_ptr.*);
+}