diff --git a/src/tools/capabilities/capability_search.zig b/src/tools/capabilities/capability_search.zig index 784f81c01..2bf5a6a66 100644 --- a/src/tools/capabilities/capability_search.zig +++ b/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, + }; +} diff --git a/src/tools/skills/skill_search.zig b/src/tools/skills/skill_search.zig index abed98070..d0d1fd462 100644 --- a/src/tools/skills/skill_search.zig +++ b/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, .{});