tool_parse: unclosed namespaced tool-call openers no longer rescan the reply

zig/src/server/tool_closers.zig

@@ -0,0 +1,40 @@
+//! Where a reply's namespaced tool-call closers ``</NAME:tool_call>`` are, found in one pass and kept by NAME.
+const std = @import("std");
+const Allocator = std.mem.Allocator;
+
+/// A byte a namespaced opener's name may hold after its first.
+pub fn nameByte(ch: u8) bool {
+    return std.ascii.isAlphanumeric(ch) or ch == '_' or ch >= 0x80 or ch == '.' or ch == '-';
+}
+
+/// One name's closers in order, and how many of them the searches have passed.
+const Seen = struct { at: std.ArrayList(usize) = .empty, passed: usize = 0 };
+
+pub const Closers = struct {
+    text: []const u8,
+    by_name: ?std.StringHashMapUnmanaged(Seen) = null, // read at the first lookup: none without a namespaced opener
+
+    /// Every ``</NAME:tool_call>`` in `text`, NAME lowered (ASCII only, as ``findIgnoreCasePos`` compares).
+    fn read(a: Allocator, text: []const u8) Allocator.Error!std.StringHashMapUnmanaged(Seen) {
+        var by_name: std.StringHashMapUnmanaged(Seen) = .empty;
+        var pos: usize = 0;
+        while (std.mem.indexOfPos(u8, text, pos, "</")) |s| {
+            pos = s + 1;
+            var j = s + 2;
+            while (j < text.len and nameByte(text[j])) j += 1;
+            if (!std.ascii.startsWithIgnoreCase(text[j..], ":tool_call>")) continue;
+            const got = try by_name.getOrPut(a, try std.ascii.allocLowerString(a, text[s + 2 .. j]));
+            if (!got.found_existing) got.value_ptr.* = .{};
+            try got.value_ptr.at.append(a, s);
+        }
+        return by_name;
+    }
+
+    /// Where ``</name:tool_call>`` (any case) first starts at or after `from`; `from` must not go down between calls.
+    pub fn next(c: *Closers, a: Allocator, name: []const u8, from: usize) Allocator.Error!?usize {
+        if (c.by_name == null) c.by_name = try read(a, c.text);
+        const seen = c.by_name.?.getPtr(try std.ascii.allocLowerString(a, name)) orelse return null;
+        while (seen.passed < seen.at.items.len and seen.at.items[seen.passed] < from) seen.passed += 1;
+        return if (seen.passed < seen.at.items.len) seen.at.items[seen.passed] else null;
+    }
+};

zig/src/server/tool_parse.zig

@@ -3,6 +3,7 @@ const std = @import("std");
 const json = @import("json");
 const ids = @import("ids.zig");
 const reply_text = @import("reply_text.zig");
+const tool_closers = @import("tool_closers.zig");
 const tool_params = @import("tool_params.zig");
 const tool_specs = @import("tool_specs.zig");
 const tool_streaming = @import("tool_streaming.zig");
@@ -64,17 +65,18 @@ fn envelopes(a: Allocator, text: []const u8) Allocator.Error![]Envelope {
         pos = e + 12;
     }
     pos = 0;
+    var closers: tool_closers.Closers = .{ .text = text };
     while (std.mem.indexOfScalarPos(u8, text, pos, '<')) |s| {
         pos = s + 1;
         var j = s + 1;
         if (j >= text.len or !(std.ascii.isAlphabetic(text[j]) or text[j] == '_')) continue;
         j += 1;
-        while (j < text.len and (word(text[j]) or text[j] == '.' or text[j] == '-')) j += 1;
+        while (j < text.len and tool_closers.nameByte(text[j])) j += 1;
         if (!std.ascii.startsWithIgnoreCase(text[j..], ":tool_call>")) continue;
-        const closer = try std.mem.concat(a, u8, &.{ "</", text[s + 1 .. j], ":tool_call>" });
-        const e = findCI(text, closer, j + 11) orelse continue;
-        try found.append(a, .{ .start = s, .end = e + closer.len, .payload = strip(text[j + 11 .. e]) });
-        pos = e + closer.len;
+        const e = (try closers.next(a, text[s + 1 .. j], j + 11)) orelse continue;
+        const end = e + j - s + 12; // ``</NAME:tool_call>``
+        try found.append(a, .{ .start = s, .end = end, .payload = strip(text[j + 11 .. e]) });
+        pos = end;
     }
     pos = 0;
     while (find(text, "<|tool_call>", pos)) |s| {

zig/src/server/tool_parse_test.zig

@@ -2,6 +2,7 @@
 const std = @import("std");
 const json = @import("json");
 const tool_parse = @import("tool_parse.zig");
+const tool_closers = @import("tool_closers.zig");
 const parse = tool_parse.parse;
 const closeCall = tool_parse.closeCall;
 const repairGlmKeys = tool_parse.repairGlmKeys;
@@ -213,6 +214,52 @@ const nested_calls = [_][]const u8{
     "<tool_call>{\"name\":\"lookup\",\"arguments\":{}",
 };
 
+test "namespaced openers without closers read in linear time" {
+    var arena: std.heap.ArenaAllocator = .init(std.testing.allocator);
+    defer arena.deinit();
+    const a = arena.allocator();
+    const tools = (try json.parse(a, "[{\"type\":\"function\",\"function\":{\"name\":\"lookup\"}}]")).ok.array;
+    // one namespace repeated, then a new namespace per opener
+    for ([_]bool{ false, true }) |distinct| {
+        var ns: [2]u64 = undefined;
+        for ([_]usize{ 1000, 8000 }, &ns) |n, *best| {
+            var reply: std.ArrayList(u8) = .empty;
+            for (0..n) |i| try reply.print(a, "<ns{d}:tool_call>", .{if (distinct) i else 0});
+            best.* = std.math.maxInt(u64);
+            for (0..3) |_| {
+                var run: std.heap.ArenaAllocator = .init(std.testing.allocator);
+                defer run.deinit();
+                const start = std.Io.Clock.awake.now(std.testing.io).nanoseconds;
+                const r = try parse(run.allocator(), reply.items, tools, null);
+                best.* = @min(best.*, @as(u64, @intCast(std.Io.Clock.awake.now(std.testing.io).nanoseconds - start)));
+                try std.testing.expect(r.calls == null);
+            }
+        }
+        // 8x the openers: near 8x the time, not 64x
+        try std.testing.expect(ns[1] <= 20 * std.time.ns_per_ms or ns[1] <= 24 * ns[0]);
+    }
+}
+
+test "a closer found by name is the one a search for it finds" {
+    var prng: std.Random.DefaultPrng = .init(0x7f4a7c15);
+    const r = prng.random();
+    const pieces = [_][]const u8{ "</a:tool_call>", "</A:TOOL_CALL>", "</ab:tool_call>", "</b.c:tool_call>", "</a:tool_cal>", "</", "<", "a", ":tool_call>", "x", "</x\u{e9}:tool_call>", "</x\u{c9}:tool_call>" };
+    const names = [_][]const u8{ "a", "A", "ab", "b.c", "B.C", "x\u{e9}", "x\u{c9}", "zz" };
+    for (0..5_000) |_| {
+        var arena: std.heap.ArenaAllocator = .init(std.testing.allocator);
+        defer arena.deinit();
+        const a = arena.allocator();
+        var text: std.ArrayList(u8) = .empty;
+        for (0..r.uintLessThan(usize, 16)) |_| try text.appendSlice(a, pieces[r.uintLessThan(usize, pieces.len)]);
+        var closers: tool_closers.Closers = .{ .text = text.items };
+        // every start in order, every name at each, as the parser asks
+        for (0..text.items.len + 1) |from| for (names) |name| {
+            const closer = try std.mem.concat(a, u8, &.{ "</", name, ":tool_call>" });
+            try std.testing.expectEqual(std.ascii.findIgnoreCasePos(text.items, from, closer), try closers.next(a, name, from));
+        };
+    }
+}
+
 test "DSML parameters contain complete and unfinished legacy calls without consuming adjacent calls" {
     var arena: std.heap.ArenaAllocator = .init(std.testing.allocator);
     defer arena.deinit();