Skip to content

Commit 39329c9

Browse files
scoootscooobscoootscooob
committed
Project Tool Search target calls in transcripts
1 parent 8f4e9c8 commit 39329c9

3 files changed

Lines changed: 277 additions & 17 deletions

File tree

‎src/agents/pi-embedded-runner/run/attempt.ts‎

Lines changed: 55 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -180,13 +180,15 @@ import {
180180
applyToolSearchCatalog,
181181
clearToolSearchCatalog,
182182
createToolSearchCatalogRef,
183+
projectToolSearchTargetTranscriptMessages,
183184
resolveToolSearchConfig,
184185
TOOL_CALL_RAW_TOOL_NAME,
185186
TOOL_DESCRIBE_RAW_TOOL_NAME,
186187
TOOL_SEARCH_CODE_MODE_TOOL_NAME,
187188
TOOL_SEARCH_RAW_TOOL_NAME,
188189
type ToolSearchCatalogRef,
189190
type ToolSearchCatalogToolExecutor,
191+
type ToolSearchTargetTranscriptProjection,
190192
} from "../../tool-search.js";
191193
import { shouldAllowProviderOwnedThinkingReplay } from "../../transcript-policy.js";
192194
import { normalizeUsage, type NormalizedUsage } from "../../usage.js";
@@ -977,6 +979,7 @@ export async function runEmbeddedAttempt(
977979
const toolSearchCatalogRef: ToolSearchCatalogRef | undefined = toolSearchControlsEnabledForRun
978980
? createToolSearchCatalogRef()
979981
: undefined;
982+
const toolSearchTargetTranscriptProjections: ToolSearchTargetTranscriptProjection[] = [];
980983
const toolsRaw = !shouldConstructTools
981984
? []
982985
: (() => {
@@ -2662,20 +2665,47 @@ export async function runEmbeddedAttempt(
26622665
getCompactionCount,
26632666
getLastCompactionTokensAfter,
26642667
} = subscription;
2665-
toolSearchCatalogExecutor = async (toolParams) =>
2666-
await runToolLifecycle({
2667-
toolName: toolParams.toolName,
2668-
toolCallId: toolParams.toolCallId,
2669-
args: toolParams.input,
2670-
execute: async () =>
2671-
await toolParams.tool.execute(
2672-
toolParams.toolCallId,
2673-
toolParams.input,
2674-
toolParams.signal ?? runAbortController.signal,
2675-
toolParams.onUpdate,
2676-
undefined as never,
2677-
),
2678-
});
2668+
toolSearchCatalogExecutor = async (toolParams) => {
2669+
try {
2670+
const result = await runToolLifecycle({
2671+
toolName: toolParams.toolName,
2672+
toolCallId: toolParams.toolCallId,
2673+
args: toolParams.input,
2674+
execute: async () =>
2675+
await toolParams.tool.execute(
2676+
toolParams.toolCallId,
2677+
toolParams.input,
2678+
toolParams.signal ?? runAbortController.signal,
2679+
toolParams.onUpdate,
2680+
undefined as never,
2681+
),
2682+
});
2683+
toolSearchTargetTranscriptProjections.push({
2684+
parentToolCallId: toolParams.parentToolCallId,
2685+
toolCallId: toolParams.toolCallId,
2686+
toolName: toolParams.toolName,
2687+
input: toolParams.input,
2688+
result,
2689+
timestamp: Date.now(),
2690+
});
2691+
return result;
2692+
} catch (error) {
2693+
const message = formatErrorMessage(error);
2694+
toolSearchTargetTranscriptProjections.push({
2695+
parentToolCallId: toolParams.parentToolCallId,
2696+
toolCallId: toolParams.toolCallId,
2697+
toolName: toolParams.toolName,
2698+
input: toolParams.input,
2699+
result: {
2700+
content: [{ type: "text", text: message }],
2701+
details: { status: "error", error: message },
2702+
},
2703+
isError: true,
2704+
timestamp: Date.now(),
2705+
});
2706+
throw error;
2707+
}
2708+
};
26792709

26802710
const queueHandle: EmbeddedPiQueueHandle & {
26812711
kind: "embedded";
@@ -3230,11 +3260,17 @@ export async function runEmbeddedAttempt(
32303260
messages: activeSession.messages,
32313261
note: `images: prompt=${imageResult.images.length}`,
32323262
});
3263+
const trajectoryProviderVisibleTools = toTrajectoryToolDefinitions(effectiveTools);
32333264
trajectoryRecorder?.recordEvent("context.compiled", {
32343265
systemPrompt: systemPromptForHook,
32353266
prompt: promptForModel,
32363267
messages: activeSession.messages,
3237-
tools: toTrajectoryToolDefinitions(effectiveTools),
3268+
tools: toTrajectoryToolDefinitions(
3269+
toolSearch.compacted ? uncompactedEffectiveTools : effectiveTools,
3270+
),
3271+
...(toolSearch.compacted
3272+
? { providerVisibleTools: trajectoryProviderVisibleTools }
3273+
: {}),
32383274
imagesCount: imageResult.images.length,
32393275
streamStrategy,
32403276
transport: effectiveAgentTransport,
@@ -3608,7 +3644,10 @@ export async function runEmbeddedAttempt(
36083644
);
36093645
}
36103646
}
3611-
messagesSnapshot = snapshotSelection.messagesSnapshot;
3647+
messagesSnapshot = projectToolSearchTargetTranscriptMessages(
3648+
snapshotSelection.messagesSnapshot,
3649+
toolSearchTargetTranscriptProjections,
3650+
);
36123651
sessionIdUsed = snapshotSelection.sessionIdUsed;
36133652

36143653
lastAssistant = messagesSnapshot

‎src/agents/tool-search.test.ts‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import {
1212
clearToolSearchCatalog,
1313
createToolSearchCatalogRef,
1414
createToolSearchTools,
15+
projectToolSearchTargetTranscriptMessages,
1516
TOOL_CALL_RAW_TOOL_NAME,
1617
TOOL_DESCRIBE_RAW_TOOL_NAME,
1718
TOOL_SEARCH_CODE_MODE_TOOL_NAME,
@@ -473,6 +474,7 @@ describe("Tool Search", () => {
473474
tool: expect.objectContaining({ name: "fake_lifecycle" }),
474475
toolName: "fake_lifecycle",
475476
toolCallId: "tool_search_code:call-lifecycle:fake_lifecycle:1",
477+
parentToolCallId: "call-lifecycle",
476478
input: { value: "ok" },
477479
signal: expect.any(AbortSignal),
478480
onUpdate,
@@ -496,13 +498,73 @@ describe("Tool Search", () => {
496498
tool: expect.objectContaining({ name: "fake_lifecycle" }),
497499
toolName: "fake_lifecycle",
498500
toolCallId: "tool_search_code:call-lifecycle-structured:fake_lifecycle:1",
501+
parentToolCallId: "call-lifecycle-structured",
499502
input: { value: "structured" },
500503
signal: abortController.signal,
501504
onUpdate,
502505
}),
503506
);
504507
});
505508

509+
it("projects target tool calls after their Tool Search wrapper result", () => {
510+
const messages = [
511+
{
512+
role: "assistant",
513+
content: [
514+
{
515+
type: "toolCall",
516+
id: "wrapper-call",
517+
name: TOOL_CALL_RAW_TOOL_NAME,
518+
arguments: { id: "fake_target", args: { value: "ok" } },
519+
},
520+
],
521+
},
522+
{
523+
role: "toolResult",
524+
toolCallId: "wrapper-call",
525+
toolName: TOOL_CALL_RAW_TOOL_NAME,
526+
content: [{ type: "text", text: "wrapped" }],
527+
},
528+
{
529+
role: "assistant",
530+
content: [{ type: "text", text: "done" }],
531+
},
532+
];
533+
534+
const projected = projectToolSearchTargetTranscriptMessages(messages as never, [
535+
{
536+
parentToolCallId: "wrapper-call",
537+
toolCallId: "tool_search_code:wrapper-call:fake_target:1",
538+
toolName: "fake_target",
539+
input: { value: "ok" },
540+
result: jsonResult({ ok: true }),
541+
timestamp: 123,
542+
},
543+
]);
544+
545+
expect(projected).toHaveLength(5);
546+
expect(projected[2]).toMatchObject({
547+
role: "assistant",
548+
content: [
549+
{
550+
type: "toolCall",
551+
id: "tool_search_code:wrapper-call:fake_target:1",
552+
name: "fake_target",
553+
arguments: { value: "ok" },
554+
input: { value: "ok" },
555+
},
556+
],
557+
});
558+
expect(projected[3]).toMatchObject({
559+
role: "toolResult",
560+
toolCallId: "tool_search_code:wrapper-call:fake_target:1",
561+
toolName: "fake_target",
562+
isError: false,
563+
content: [{ type: "text", text: JSON.stringify({ ok: true }, null, 2) }],
564+
});
565+
expect(projected[4]).toBe(messages[2]);
566+
});
567+
506568
it("does not execute fire-and-forget bridged calls after code returns", async () => {
507569
const codeTool = fakeTool(TOOL_SEARCH_CODE_MODE_TOOL_NAME, "code mode");
508570
const target = pluginTool("fake_fire_and_forget", "Should not run unless awaited");

‎src/agents/tool-search.ts‎

Lines changed: 160 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
import { spawn } from "node:child_process";
22
import os from "node:os";
3-
import type { AgentToolResult, AgentToolUpdateCallback } from "@mariozechner/pi-agent-core";
3+
import type {
4+
AgentMessage,
5+
AgentToolResult,
6+
AgentToolUpdateCallback,
7+
} from "@mariozechner/pi-agent-core";
48
import type { ToolDefinition } from "@mariozechner/pi-coding-agent";
59
import { Type } from "typebox";
610
import type { OpenClawConfig } from "../config/types.openclaw.js";
@@ -37,11 +41,22 @@ export type ToolSearchCatalogToolExecutor = (params: {
3741
tool: CatalogTool;
3842
toolName: string;
3943
toolCallId: string;
44+
parentToolCallId?: string;
4045
input: unknown;
4146
signal?: AbortSignal;
4247
onUpdate?: AgentToolUpdateCallback<unknown>;
4348
}) => Promise<AgentToolResult<unknown>>;
4449

50+
export type ToolSearchTargetTranscriptProjection = {
51+
parentToolCallId?: string;
52+
toolCallId: string;
53+
toolName: string;
54+
input: unknown;
55+
result?: unknown;
56+
isError?: boolean;
57+
timestamp?: number;
58+
};
59+
4560
export type ToolSearchConfig = {
4661
enabled: boolean;
4762
mode: ToolSearchMode;
@@ -518,6 +533,149 @@ function dropToolSearchControlTools(tools: AnyAgentTool[]): AnyAgentTool[] {
518533
return tools.filter((tool) => !TOOL_SEARCH_CONTROL_TOOL_NAMES.has(tool.name));
519534
}
520535

536+
function readMessageToolResultId(message: AgentMessage): string | undefined {
537+
const record = message as unknown as Record<string, unknown>;
538+
const role = typeof record.role === "string" ? record.role : "";
539+
const canUseDirectId = role === "toolResult" || role === "tool";
540+
const direct = record.toolCallId ?? record.toolUseId ?? record.tool_use_id;
541+
if (canUseDirectId && typeof direct === "string" && direct.trim()) {
542+
return direct;
543+
}
544+
const content = record.content;
545+
if (!Array.isArray(content)) {
546+
return undefined;
547+
}
548+
for (const block of content) {
549+
if (!isRecord(block)) {
550+
continue;
551+
}
552+
if (block.type !== "toolResult") {
553+
continue;
554+
}
555+
const nested = block.toolCallId ?? block.toolUseId ?? block.tool_use_id ?? block.id;
556+
if (typeof nested === "string" && nested.trim()) {
557+
return nested;
558+
}
559+
}
560+
return undefined;
561+
}
562+
563+
function textFromToolSearchProjectionResult(result: unknown, isError: boolean): string {
564+
if (isRecord(result)) {
565+
const details = isRecord(result.details) ? result.details : undefined;
566+
const detailError = details?.error;
567+
if (typeof detailError === "string" && detailError.trim()) {
568+
return detailError;
569+
}
570+
const content = result.content;
571+
if (Array.isArray(content)) {
572+
const text = content
573+
.map((item) => (isRecord(item) && typeof item.text === "string" ? item.text : ""))
574+
.filter(Boolean)
575+
.join("\n");
576+
if (text.trim()) {
577+
return text;
578+
}
579+
}
580+
}
581+
const safe = toJsonSafe(result);
582+
if (typeof safe === "string") {
583+
return safe;
584+
}
585+
const encoded = JSON.stringify(safe);
586+
if (typeof encoded === "string") {
587+
return encoded;
588+
}
589+
return isError ? "Tool Search target tool failed." : "Tool Search target tool completed.";
590+
}
591+
592+
function buildToolSearchTargetTranscriptMessages(
593+
projection: ToolSearchTargetTranscriptProjection,
594+
): AgentMessage[] {
595+
const input = toJsonSafe(projection.input);
596+
const timestamp = projection.timestamp ?? Date.now();
597+
const resultRecord = isRecord(projection.result) ? projection.result : undefined;
598+
const resultContent =
599+
Array.isArray(resultRecord?.content) && resultRecord.content.length > 0
600+
? toJsonSafe(resultRecord.content)
601+
: [
602+
{
603+
type: "text",
604+
text: textFromToolSearchProjectionResult(
605+
projection.result,
606+
projection.isError === true,
607+
),
608+
},
609+
];
610+
return [
611+
{
612+
role: "assistant",
613+
content: [
614+
{
615+
type: "toolCall",
616+
id: projection.toolCallId,
617+
name: projection.toolName,
618+
arguments: input,
619+
input,
620+
},
621+
],
622+
stopReason: "toolUse",
623+
timestamp,
624+
} as unknown as AgentMessage,
625+
{
626+
role: "toolResult",
627+
toolCallId: projection.toolCallId,
628+
toolName: projection.toolName,
629+
isError: projection.isError === true,
630+
content: resultContent,
631+
timestamp,
632+
} as unknown as AgentMessage,
633+
];
634+
}
635+
636+
export function projectToolSearchTargetTranscriptMessages(
637+
messages: AgentMessage[],
638+
projections: readonly ToolSearchTargetTranscriptProjection[],
639+
): AgentMessage[] {
640+
if (projections.length === 0) {
641+
return messages;
642+
}
643+
const byParent = new Map<string, ToolSearchTargetTranscriptProjection[]>();
644+
const unmatched: ToolSearchTargetTranscriptProjection[] = [];
645+
for (const projection of projections) {
646+
const parent = projection.parentToolCallId?.trim();
647+
if (!parent) {
648+
unmatched.push(projection);
649+
continue;
650+
}
651+
const group = byParent.get(parent) ?? [];
652+
group.push(projection);
653+
byParent.set(parent, group);
654+
}
655+
const inserted = new Set<ToolSearchTargetTranscriptProjection>();
656+
const projected: AgentMessage[] = [];
657+
for (const message of messages) {
658+
projected.push(message);
659+
const toolResultId = readMessageToolResultId(message);
660+
const group = toolResultId ? byParent.get(toolResultId) : undefined;
661+
if (!group) {
662+
continue;
663+
}
664+
for (const projection of group) {
665+
projected.push(...buildToolSearchTargetTranscriptMessages(projection));
666+
inserted.add(projection);
667+
}
668+
}
669+
for (const projection of [...unmatched, ...projections]) {
670+
if (inserted.has(projection)) {
671+
continue;
672+
}
673+
projected.push(...buildToolSearchTargetTranscriptMessages(projection));
674+
inserted.add(projection);
675+
}
676+
return projected;
677+
}
678+
521679
export function createToolSearchCatalogRef(): ToolSearchCatalogRef {
522680
return {};
523681
}
@@ -888,6 +1046,7 @@ class ToolSearchRuntime {
8881046
tool: entry.tool,
8891047
toolName: entry.name,
8901048
toolCallId,
1049+
parentToolCallId: options?.parentToolCallId,
8911050
input: input ?? {},
8921051
signal: options?.signal ?? this.ctx.abortSignal,
8931052
onUpdate: options?.onUpdate,

0 commit comments

Comments
 (0)