Merge pull request #1055 from vercel-labs/fazxes/streamline-capability-search

Streamline capability search result assembly

src/tools/capabilities/capability_search.zig

@@ -30,12 +30,6 @@ const Input = struct {
     }
 };
 
-const ProjectionCheck = union(enum) {
-    valid,
-    invalid,
-    authentication_identity_changed,
-};
-
 pub fn decode(
     ctx: tool_dispatch.DispatchContext,
     args_json: []const u8,
@@ -125,17 +119,14 @@ pub fn call(
         (output_cap - 512) / 2
     else
         output_cap;
-    const skill_output = if (searches_skills)
-        skill_search.searchRequest(ctx, request, domain_cap) catch |err| switch (err) {
+    var skill_result: ?skill_search.SearchResult = null;
+    if (searches_skills) {
+        skill_result = skill_search.searchRequest(ctx, request, domain_cap) catch |err| switch (err) {
             error.OutOfMemory => return error.OutOfMemory,
             else => return executionFailure(ctx.allocator, "skill", err),
-        }
-    else
-        try ctx.allocator.dupe(
-            u8,
-            "{\"skills\":[],\"count\":0,\"total_matches\":0,\"more_available\":false,\"next_cursor\":null}",
-        );
-    defer ctx.allocator.free(skill_output);
+        };
+    }
+    defer if (skill_result) |*result| result.deinit(ctx.allocator);
 
     var mcp_result: tool_mcp_runtime.SearchResult = if (searches_mcp)
         try searchMcp(ctx, request, domain_cap)
@@ -146,7 +137,7 @@ pub fn call(
 
     const combined = combineProjected(
         ctx.allocator,
-        skill_output,
+        if (skill_result) |*result| result else null,
         mcp_result.model_output,
         output_cap,
     ) catch |err| switch (err) {
@@ -218,100 +209,45 @@ fn searchMcp(
 
 fn combineProjected(
     alloc: Allocator,
-    skill_output: []const u8,
+    skill_result: ?*const skill_search.SearchResult,
     mcp_output: []const u8,
     max_bytes: usize,
 ) ![]u8 {
-    var skills_parsed = try std.json.parseFromSlice(std.json.Value, alloc, skill_output, .{});
-    defer skills_parsed.deinit();
     var mcp_parsed = try std.json.parseFromSlice(std.json.Value, alloc, mcp_output, .{});
     defer mcp_parsed.deinit();
-    if (skills_parsed.value != .object or mcp_parsed.value != .object) {
-        return error.InvalidCapabilitySearchResult;
-    }
-    const skills_value = skills_parsed.value.object.get("skills") orelse
-        return error.InvalidCapabilitySearchResult;
+    if (mcp_parsed.value != .object) return error.InvalidCapabilitySearchResult;
     const mcp_tools_value = mcp_parsed.value.object.get("tools") orelse
         return error.InvalidCapabilitySearchResult;
-    if (skills_value != .array or mcp_tools_value != .array) {
-        return error.InvalidCapabilitySearchResult;
-    }
-
-    const skill_items = skills_value.array.items;
+    if (mcp_tools_value != .array) return error.InvalidCapabilitySearchResult;
     const mcp_items = mcp_tools_value.array.items;
-    const skill_count = skill_items.len;
-    const mcp_count = mcp_items.len;
-    var include_authentication = mcp_parsed.value.object.get("authentication_required") != null;
 
-    while (true) {
-        const raw = renderCombined(
-            alloc,
-            skill_items[0..skill_count],
-            mcp_items[0..mcp_count],
-            objectNonNegativeCount(skills_parsed.value, "total_matches"),
-            objectNonNegativeCount(mcp_parsed.value, "total_matches"),
-            skills_parsed.value.object.get("state"),
-            if (include_authentication)
-                mcp_parsed.value.object.get("authentication_required")
-            else
-                null,
-            mcp_parsed.value.object.get("state"),
-            mcp_parsed.value.object.get("error"),
-            mcp_parsed.value.object.get("context_limit"),
-        ) catch |err| switch (err) {
-            error.WriteFailed => return error.OutOfMemory,
-            else => return err,
-        };
-        defer alloc.free(raw);
-
-        const projected = @constCast(try tool_result_limits.prepareModelOutput(
-            alloc,
-            "capability_search",
-            raw,
-            max_bytes,
-        ));
-        errdefer alloc.free(projected);
-        const check = try checkProjection(
-            alloc,
-            projected,
-            skill_items[0..skill_count],
-            mcp_items[0..mcp_count],
-            if (include_authentication)
-                mcp_parsed.value.object.get("authentication_required")
-            else
-                null,
-        );
-        switch (check) {
-            .valid => {
-                const second = try tool_result_limits.prepareModelOutput(
-                    alloc,
-                    "capability_search",
-                    projected,
-                    max_bytes,
-                );
-                defer alloc.free(@constCast(second));
-                if (std.mem.eql(u8, projected, second)) return projected;
-                alloc.free(projected);
-                return error.CapabilitySearchResultLimitTooSmall;
-            },
-            .authentication_identity_changed => include_authentication = false,
-            .invalid,
-            => {
-                alloc.free(projected);
-                return error.CapabilitySearchResultLimitTooSmall;
-            },
-        }
-        alloc.free(projected);
-    }
+    const combined = renderCombined(
+        alloc,
+        if (skill_result) |result| result.itemsJson() else "",
+        if (skill_result) |result| result.count else 0,
+        mcp_items,
+        if (skill_result) |result| result.total_matches else 0,
+        objectNonNegativeCount(mcp_parsed.value, "total_matches"),
+        mcp_parsed.value.object.get("authentication_required"),
+        mcp_parsed.value.object.get("state"),
+        mcp_parsed.value.object.get("error"),
+        mcp_parsed.value.object.get("context_limit"),
+    ) catch |err| switch (err) {
+        error.WriteFailed => return error.OutOfMemory,
+        else => return err,
+    };
+    if (combined.len <= max_bytes) return combined;
+    alloc.free(combined);
+    return error.CapabilitySearchResultLimitTooSmall;
 }
 
 fn renderCombined(
     alloc: Allocator,
-    skills: []const std.json.Value,
+    skills_json: []const u8,
+    skill_count: usize,
     mcp_tools: []const std.json.Value,
     skill_total_matches: usize,
     mcp_total_matches: usize,
-    skill_state: ?std.json.Value,
     authentication_required: ?std.json.Value,
     mcp_state: ?std.json.Value,
     mcp_error: ?std.json.Value,
@@ -320,10 +256,7 @@ fn renderCombined(
     var out: std.Io.Writer.Allocating = .init(alloc);
     defer out.deinit();
     try out.writer.writeAll("{\"skills\":[");
-    for (skills, 0..) |value, index| {
-        if (index > 0) try out.writer.writeByte(',');
-        try std.json.Stringify.value(value, .{}, &out.writer);
-    }
+    try out.writer.writeAll(skills_json);
     try out.writer.writeAll("],\"mcp_tools\":[");
     for (mcp_tools, 0..) |value, index| {
         if (index > 0) try out.writer.writeByte(',');
@@ -332,7 +265,7 @@ fn renderCombined(
     try out.writer.print(
         "],\"counts\":{{\"skills\":{d},\"mcp_tools\":{d}}},\"total_matches\":{{\"skills\":{d},\"mcp_tools\":{d}}}",
         .{
-            skills.len,
+            skill_count,
             mcp_tools.len,
             skill_total_matches,
             mcp_total_matches,
@@ -340,17 +273,12 @@ fn renderCombined(
     );
     if (skill_total_matches == 0 and
         mcp_total_matches == 0 and
-        skill_state == null and
         authentication_required == null and
         mcp_state == null and
         mcp_error == null)
     {
         try out.writer.writeAll(",\"state\":\"no_match\"");
     }
-    if (skill_state) |value| {
-        try out.writer.writeAll(",\"skill_state\":");
-        try std.json.Stringify.value(value, .{}, &out.writer);
-    }
     if (authentication_required) |value| {
         try out.writer.writeAll(",\"authentication_required\":");
         try std.json.Stringify.value(value, .{}, &out.writer);
@@ -371,70 +299,6 @@ fn renderCombined(
     return out.toOwnedSlice();
 }
 
-fn checkProjection(
-    alloc: Allocator,
-    projected: []const u8,
-    skills: []const std.json.Value,
-    mcp_tools: []const std.json.Value,
-    authentication_required: ?std.json.Value,
-) !ProjectionCheck {
-    var parsed = std.json.parseFromSlice(std.json.Value, alloc, projected, .{}) catch |err| switch (err) {
-        error.OutOfMemory => return error.OutOfMemory,
-        else => return .invalid,
-    };
-    defer parsed.deinit();
-    if (parsed.value != .object) return .invalid;
-    const projected_skills = parsed.value.object.get("skills") orelse return .invalid;
-    const projected_mcp = parsed.value.object.get("mcp_tools") orelse return .invalid;
-    if (projected_skills != .array or projected_mcp != .array or
-        projected_skills.array.items.len != skills.len or
-        projected_mcp.array.items.len != mcp_tools.len)
-    {
-        return .invalid;
-    }
-
-    for (projected_skills.array.items, skills) |projected_skill, skill| {
-        if (!objectStringFieldsEqual(projected_skill, skill, "name", "location")) {
-            return .invalid;
-        }
-    }
-    for (projected_mcp.array.items, mcp_tools) |projected_tool, tool| {
-        if (!objectStringFieldsEqual(projected_tool, tool, "name", "server")) {
-            return .invalid;
-        }
-    }
-    if (authentication_required) |expected| {
-        const projected_auth = parsed.value.object.get("authentication_required") orelse
-            return .authentication_identity_changed;
-        if (!objectStringFieldEqual(projected_auth, expected, "server")) {
-            return .authentication_identity_changed;
-        }
-    }
-    return .valid;
-}
-
-fn objectStringFieldsEqual(
-    actual: std.json.Value,
-    expected: std.json.Value,
-    first: []const u8,
-    second: []const u8,
-) bool {
-    return objectStringFieldEqual(actual, expected, first) and
-        objectStringFieldEqual(actual, expected, second);
-}
-
-fn objectStringFieldEqual(
-    actual: std.json.Value,
-    expected: std.json.Value,
-    field: []const u8,
-) bool {
-    if (actual != .object or expected != .object) return false;
-    const actual_value = actual.object.get(field) orelse return false;
-    const expected_value = expected.object.get(field) orelse return false;
-    return actual_value == .string and expected_value == .string and
-        std.mem.eql(u8, actual_value.string, expected_value.string);
-}
-
 fn objectNonNegativeCount(value: std.json.Value, field: []const u8) usize {
     const candidate = value.object.get(field) orelse return 0;
     if (candidate != .integer or candidate.integer < 0) return 0;
@@ -515,7 +379,8 @@ test "capability search combines bounded skill and MCP results" {
         "{\"skills\":[{\"name\":\"mail-helper\",\"description\":\"Send email\",\"location\":\"/skills/mail-helper\"}],\"count\":1,\"total_matches\":2,\"more_available\":true,\"next_cursor\":\"c1:s:1:1:1\"}";
     const mcp =
         "{\"tools\":[{\"name\":\"mcp_mail_send\",\"server\":\"mail\",\"description\":\"Send email\"}],\"count\":1,\"total_matches\":3,\"more_available\":true,\"next_cursor\":\"c1:m:1:1:1\",\"authentication_required\":{\"server\":\"TOKEN=runtime-auth-secret\",\"message\":\"authenticate\"}}";
-    const output = try combineProjected(alloc, skills, mcp, 4096);
+    const skill_result = testSkillResult(skills, 1, 2);
+    const output = try combineProjected(alloc, &skill_result, mcp, 4096);
     defer alloc.free(output);
 
     var parsed = try std.json.parseFromSlice(std.json.Value, alloc, output, .{});
@@ -530,6 +395,20 @@ test "capability search combines bounded skill and MCP results" {
     try std.testing.expect(parsed.value.object.get("more_available") == null);
     try std.testing.expect(parsed.value.object.get("next_cursors") == null);
     try std.testing.expect(parsed.value.object.get("authentication_required") != null);
+
+    const projected_again = try tool_result_limits.prepareModelOutput(
+        alloc,
+        "capability_search",
+        output,
+        output.len,
+    );
+    defer alloc.free(@constCast(projected_again));
+    try std.testing.expectEqualStrings(output, projected_again);
+
+    try std.testing.expectError(
+        error.CapabilitySearchResultLimitTooSmall,
+        combineProjected(alloc, &skill_result, mcp, output.len - 1),
+    );
 }
 
 test "capability search combined projection releases every allocation failure" {
@@ -539,7 +418,8 @@ test "capability search combined projection releases every allocation failure" {
                 "{\"skills\":[{\"name\":\"mail-helper\",\"description\":\"Send email\",\"location\":\"/skills/mail-helper\"}],\"count\":1,\"total_matches\":1,\"more_available\":false,\"next_cursor\":null}";
             const mcp =
                 "{\"tools\":[{\"name\":\"mcp_mail_send\",\"server\":\"mail\",\"description\":\"Send email\"}],\"count\":1,\"total_matches\":1,\"more_available\":false,\"next_cursor\":null}";
-            const output = try combineProjected(alloc, skills, mcp, 1024);
+            const skill_result = testSkillResult(skills, 1, 1);
+            const output = try combineProjected(alloc, &skill_result, mcp, 1024);
             alloc.free(output);
         }
     };
@@ -552,7 +432,8 @@ test "capability search marks an empty search as terminal" {
         "{\"skills\":[],\"count\":0,\"total_matches\":0,\"more_available\":false,\"next_cursor\":null}";
     const mcp =
         "{\"tools\":[],\"count\":0,\"total_matches\":0,\"more_available\":false,\"next_cursor\":null}";
-    const output = try combineProjected(alloc, skills, mcp, 4096);
+    const skill_result = testSkillResult(skills, 0, 0);
+    const output = try combineProjected(alloc, &skill_result, mcp, 4096);
     defer alloc.free(output);
 
     var parsed = try std.json.parseFromSlice(std.json.Value, alloc, output, .{});
@@ -563,3 +444,15 @@ test "capability search marks an empty search as terminal" {
     );
     try std.testing.expectEqual(@as(usize, 5), parsed.value.object.count());
 }
+
+fn testSkillResult(output: []const u8, count: usize, total_matches: usize) skill_search.SearchResult {
+    const items_start = "{\"skills\":[".len;
+    const items_end = std.mem.find(u8, output, "],\"count\":") orelse unreachable;
+    return .{
+        .model_output = @constCast(output),
+        .items_start = items_start,
+        .items_end = items_end,
+        .count = count,
+        .total_matches = total_matches,
+    };
+}

src/tools/skills/skill_search.zig

@@ -9,17 +9,34 @@ const tool_result_limits = @import("../../core/tooling/tool_result_limits.zig");
 
 const Allocator = std.mem.Allocator;
 
-const ProjectionCheck = union(enum) {
-    valid,
-    invalid,
-    identity_changed: usize,
+pub const SearchResult = struct {
+    model_output: []u8,
+    items_start: usize,
+    items_end: usize,
+    count: usize,
+    total_matches: usize,
+
+    pub fn deinit(self: *SearchResult, alloc: Allocator) void {
+        alloc.free(self.model_output);
+        self.* = undefined;
+    }
+
+    pub fn itemsJson(self: *const SearchResult) []const u8 {
+        return self.model_output[self.items_start..self.items_end];
+    }
+};
+
+const RenderedSearch = struct {
+    bytes: []u8,
+    items_start: usize,
+    items_end: usize,
 };
 
 pub fn searchRequest(
     ctx: tool_dispatch.DispatchContext,
     request: capability_retrieval.Request,
     max_bytes: usize,
-) ![]u8 {
+) !SearchResult {
     var discovery = try builtin_skills.loadVisibleSkillsForTool(
         ctx.allocator,
         ctx.workspace_root,
@@ -58,7 +75,7 @@ fn renderProjectedSearch(
     skills: []const skill_runtime.Skill,
     description_limit: context_limits.Resolved,
     max_bytes: usize,
-) ![]u8 {
+) !SearchResult {
     const documents = try alloc.alloc(capability_retrieval.Document, skills.len);
     defer alloc.free(documents);
     const document_skills = try alloc.alloc(*const skill_runtime.Skill, skills.len);
@@ -105,41 +122,25 @@ fn renderProjectedSearch(
             error.WriteFailed => return error.OutOfMemory,
             else => return err,
         };
-        defer alloc.free(raw);
+        defer alloc.free(raw.bytes);
 
         const projected = @constCast(try tool_result_limits.prepareModelOutput(
             alloc,
             "capability_search",
-            raw,
+            raw.bytes,
             max_bytes,
         ));
         errdefer alloc.free(projected);
-        const check = try checkProjection(
-            alloc,
-            projected,
-            page.matches[0..retained_count],
-            document_skills[0..document_count],
-            page.total_matches,
-            next_cursor,
-        );
-        switch (check) {
-            .valid => {
-                const second = try tool_result_limits.prepareModelOutput(
-                    alloc,
-                    "capability_search",
-                    projected,
-                    max_bytes,
-                );
-                defer alloc.free(@constCast(second));
-                if (std.mem.eql(u8, projected, second)) return projected;
-            },
-            .invalid, .identity_changed => {},
-        }
-        alloc.free(projected);
-        retained_count = switch (check) {
-            .identity_changed => |index| index,
-            .valid, .invalid => if (retained_count > 0) retained_count - 1 else return error.SkillSearchResultLimitTooSmall,
+        if (std.mem.eql(u8, raw.bytes, projected)) return .{
+            .model_output = projected,
+            .items_start = raw.items_start,
+            .items_end = raw.items_end,
+            .count = retained_count,
+            .total_matches = page.total_matches,
         };
+        alloc.free(projected);
+        if (retained_count == 0) return error.SkillSearchResultLimitTooSmall;
+        retained_count -= 1;
     }
 }
 
@@ -150,10 +151,11 @@ fn renderRawSearch(
     description_limit: usize,
     total_matches: usize,
     next_cursor: ?[]const u8,
-) ![]u8 {
+) !RenderedSearch {
     var out: std.Io.Writer.Allocating = .init(alloc);
     defer out.deinit();
     try out.writer.writeAll("{\"skills\":[");
+    const items_start = out.written().len;
     for (matches, 0..) |match, index| {
         if (index > 0) try out.writer.writeByte(',');
         const skill = skills[match.document_index];
@@ -169,6 +171,7 @@ fn renderRawSearch(
         try std.json.Stringify.value(skill.path, .{}, &out.writer);
         try out.writer.writeByte('}');
     }
+    const items_end = out.written().len;
     try out.writer.print("],\"count\":{d},\"total_matches\":{d},\"more_available\":{s},\"next_cursor\":", .{
         matches.len,
         total_matches,
@@ -176,63 +179,11 @@ fn renderRawSearch(
     });
     try std.json.Stringify.value(next_cursor, .{}, &out.writer);
     try out.writer.writeByte('}');
-    return try out.toOwnedSlice();
-}
-
-fn checkProjection(
-    alloc: Allocator,
-    projected: []const u8,
-    matches: []const capability_retrieval.Match,
-    skills: []const *const skill_runtime.Skill,
-    total_matches: usize,
-    next_cursor: ?[]const u8,
-) !ProjectionCheck {
-    var parsed = std.json.parseFromSlice(std.json.Value, alloc, projected, .{}) catch |err| switch (err) {
-        error.OutOfMemory => return error.OutOfMemory,
-        else => return .invalid,
+    return .{
+        .bytes = try out.toOwnedSlice(),
+        .items_start = items_start,
+        .items_end = items_end,
     };
-    defer parsed.deinit();
-    if (parsed.value != .object) return .invalid;
-
-    const skills_value = parsed.value.object.get("skills") orelse return .invalid;
-    const count_value = parsed.value.object.get("count") orelse return .invalid;
-    const total_value = parsed.value.object.get("total_matches") orelse return .invalid;
-    const more_value = parsed.value.object.get("more_available") orelse return .invalid;
-    const cursor_value = parsed.value.object.get("next_cursor") orelse return .invalid;
-    if (skills_value != .array or
-        count_value != .integer or
-        count_value.integer < 0 or
-        total_value != .integer or
-        total_value.integer < 0 or
-        more_value != .bool or
-        more_value.bool != (next_cursor != null) or
-        skills_value.array.items.len != matches.len)
-    {
-        return .invalid;
-    }
-    const count = std.math.cast(usize, count_value.integer) orelse return .invalid;
-    const total = std.math.cast(usize, total_value.integer) orelse return .invalid;
-    if (count != matches.len or total != total_matches) return .invalid;
-    if (next_cursor) |expected| {
-        if (cursor_value != .string or !std.mem.eql(u8, cursor_value.string, expected)) {
-            return .invalid;
-        }
-    } else if (cursor_value != .null) return .invalid;
-
-    for (skills_value.array.items, matches, 0..) |value, match, index| {
-        if (value != .object) return .invalid;
-        const name = value.object.get("name") orelse return .invalid;
-        const description = value.object.get("description") orelse return .invalid;
-        const location = value.object.get("location") orelse return .invalid;
-        if (name != .string or description != .string or location != .string) return .invalid;
-        const skill = skills[match.document_index];
-        if (!std.mem.eql(u8, name.string, skill.name) or
-            !std.mem.eql(u8, location.string, skill.path))
-        {
-            return .{ .identity_changed = index };
-        }
-    }
-    return .valid;
 }
 
 test "skill search ranks metadata and returns final-projection-stable JSON" {
@@ -242,14 +193,15 @@ test "skill search ranks metadata and returns final-projection-stable JSON" {
         .{ .name = "fx-review", .description = "Review fx runtime changes", .path = "/skills/fx-review/SKILL.md", .source = .workspace_fx },
     };
     const query = try lexical_relevance.prepare("review fx runtime public");
-    const output = try renderProjectedSearch(
+    var result = try renderProjectedSearch(
         alloc,
         .{ .query = &query, .kind = .skill, .limit = capability_retrieval.default_limit },
         &skills,
         (context_limits.Values{}).skill_description_bytes,
         4096,
     );
-    defer alloc.free(output);
+    defer result.deinit(alloc);
+    const output = result.model_output;
     try std.testing.expect(std.mem.find(u8, output, "\"name\":\"fx-review\"") != null);
     try std.testing.expect(std.mem.find(u8, output, "\"count\":1") != null);
 
@@ -265,14 +217,15 @@ test "skill search preserves projected identities and verbatim descriptions" {
         .{ .name = "safe", .description = "API_KEY=description-secret-value", .path = "/skills/safe/SKILL.md", .source = .workspace_fx },
     };
     const query = try lexical_relevance.prepare("");
-    const output = try renderProjectedSearch(
+    var result = try renderProjectedSearch(
         alloc,
         .{ .query = &query, .kind = .skill, .limit = capability_retrieval.default_limit },
         &skills,
         (context_limits.Values{}).skill_description_bytes,
         4096,
     );
-    defer alloc.free(output);
+    defer result.deinit(alloc);
+    const output = result.model_output;
     try std.testing.expect(std.mem.find(u8, output, "unsafe-location") != null);
     try std.testing.expect(std.mem.find(u8, output, "\"name\":\"safe\"") != null);
     try std.testing.expect(std.mem.find(u8, output, "API_KEY=description-secret-value") != null);
@@ -294,14 +247,15 @@ test "skill search caps ranked entries and atomically omits byte overflow" {
         .{ .name = "nine", .description = "nine", .path = "/skills/nine/SKILL.md", .source = .workspace_fx },
     };
     const query = try lexical_relevance.prepare("");
-    const output = try renderProjectedSearch(
+    var result = try renderProjectedSearch(
         alloc,
         .{ .query = &query, .kind = .skill, .limit = capability_retrieval.default_limit },
         &skills,
         (context_limits.Values{}).skill_description_bytes,
         1024,
     );
-    defer alloc.free(output);
+    defer result.deinit(alloc);
+    const output = result.model_output;
 
     var parsed = try std.json.parseFromSlice(std.json.Value, alloc, output, .{});
     defer parsed.deinit();
@@ -327,14 +281,14 @@ test "skill search projection releases every allocation failure" {
                 .{ .name = "nine", .description = "nine", .path = "/skills/nine/SKILL.md", .source = .workspace_fx },
             };
             const query = try lexical_relevance.prepare("");
-            const output = try renderProjectedSearch(
+            var result = try renderProjectedSearch(
                 alloc,
                 .{ .query = &query, .kind = .skill, .limit = capability_retrieval.default_limit },
                 &skills,
                 (context_limits.Values{}).skill_description_bytes,
                 1024,
             );
-            alloc.free(output);
+            result.deinit(alloc);
         }
     };
     try std.testing.checkAllAllocationFailures(std.testing.allocator, Case.run, .{});