ShockVPN/shockbot
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: update retry logic

Browse Source
This commit is contained in:
wolfy
2026-06-02 19:53:21 -05:00
parent eb8871d6d2
commit 1de60d74fd
1 changed files with 273 additions and 273 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agents/ollama.ts
+273 -273
View File
@@ -9,64 +9,64 @@ const DEFAULT_MODEL = "qwen3.6:35b";
const MAX_ITERATIONS = 100; const MAX_ITERATIONS = 100;
interface OllamaTool { interface OllamaTool {
type: "function"; type: "function";
function: { function: {
name: string; name: string;
description: string; description: string;
parameters: Record<string, unknown>; parameters: Record<string, unknown>;
}; };
} }
async function buildMcpClient(mcpServerUrl: string): Promise<Client> { async function buildMcpClient(mcpServerUrl: string): Promise<Client> {
const client = new Client( const client = new Client(
{ name: "shockbot-agent", version: "0.1.0" }, { name: "shockbot-agent", version: "0.1.0" },
{ capabilities: {} }, { capabilities: {} },
); );
const transport = new StreamableHTTPClientTransport(new URL(mcpServerUrl)); const transport = new StreamableHTTPClientTransport(new URL(mcpServerUrl));
await client.connect(transport); await client.connect(transport);
return client; return client;
} }
async function getOllamaTools(mcpClient: Client): Promise<OllamaTool[]> { async function getOllamaTools(mcpClient: Client): Promise<OllamaTool[]> {
const { tools } = await mcpClient.listTools(); const { tools } = await mcpClient.listTools();
return tools.map((t) => ({ return tools.map((t) => ({
type: "function" as const, type: "function" as const,
function: { function: {
name: t.name, name: t.name,
description: t.description ?? "", description: t.description ?? "",
parameters: (t.inputSchema as Record<string, unknown>) ?? { parameters: (t.inputSchema as Record<string, unknown>) ?? {
type: "object", type: "object",
properties: {}, properties: {},
}, },
}, },
})); }));
} }
async function callMcpTool( async function callMcpTool(
mcpClient: Client, mcpClient: Client,
toolName: string, toolName: string,
args: Record<string, unknown>, args: Record<string, unknown>,
): Promise<string> { ): Promise<string> {
try { try {
const result = await mcpClient.callTool({ const result = await mcpClient.callTool({
name: toolName, name: toolName,
arguments: args, arguments: args,
}); });
const content = result.content as const content = result.content as
| Array<{ type: string; text?: string }> | Array<{ type: string; text?: string }>
| undefined; | undefined;
if (!content || content.length === 0) if (!content || content.length === 0)
return JSON.stringify({ success: true }); return JSON.stringify({ success: true });
const text = content const text = content
.map((c) => (c.type === "text" ? (c.text ?? "") : "")) .map((c) => (c.type === "text" ? (c.text ?? "") : ""))
.filter(Boolean) .filter(Boolean)
.join("\n"); .join("\n");
return text || JSON.stringify(result); return text || JSON.stringify(result);
} catch (err) { } catch (err) {
const msg = err instanceof Error ? err.message : String(err); const msg = err instanceof Error ? err.message : String(err);
log.debug(`Tool ${toolName} error: ${msg}`); log.debug(`Tool ${toolName} error: ${msg}`);
return JSON.stringify({ error: msg }); return JSON.stringify({ error: msg });
} }
} }
/** /**
@@ -76,260 +76,260 @@ async function callMcpTool(
* Never touches system/user/assistant messages — only tool messages. * Never touches system/user/assistant messages — only tool messages.
*/ */
function pruneToolMessages(messages: Message[], keepRecent = 6): Message[] { function pruneToolMessages(messages: Message[], keepRecent = 6): Message[] {
const toolIndices: number[] = []; const toolIndices: number[] = [];
for (let i = 0; i < messages.length; i++) { for (let i = 0; i < messages.length; i++) {
if (messages[i].role === "tool") toolIndices.push(i); if (messages[i].role === "tool") toolIndices.push(i);
} }
const pruneCount = Math.max(0, toolIndices.length - keepRecent); const pruneCount = Math.max(0, toolIndices.length - keepRecent);
if (pruneCount === 0) return messages; if (pruneCount === 0) return messages;
const toPrune = new Set(toolIndices.slice(0, pruneCount)); const toPrune = new Set(toolIndices.slice(0, pruneCount));
let pruned = 0; let pruned = 0;
const result = messages.map((msg, i) => { const result = messages.map((msg, i) => {
if (!toPrune.has(i)) return msg; if (!toPrune.has(i)) return msg;
const originalLen = const originalLen =
typeof msg.content === "string" ? msg.content.length : 0; typeof msg.content === "string" ? msg.content.length : 0;
pruned++; pruned++;
return { return {
...msg, ...msg,
content: `[pruned: was ${originalLen} chars — context limit approached]`, content: `[pruned: was ${originalLen} chars — context limit approached]`,
}; };
}); });
log.info(`» pruned ${pruned} old tool message(s) to reduce context`); log.info(`» pruned ${pruned} old tool message(s) to reduce context`);
return result; return result;
} }
async function unloadModel(ollama: Ollama, model: string): Promise<void> { async function unloadModel(ollama: Ollama, model: string): Promise<void> {
try { try {
await ollama.generate({ model, keep_alive: 0, prompt: "" }); await ollama.generate({ model, keep_alive: 0, prompt: "" });
log.info(`» unloaded model ${model} from Ollama`); log.info(`» unloaded model ${model} from Ollama`);
} catch (err) { } catch (err) {
const msg = err instanceof Error ? err.message : String(err); const msg = err instanceof Error ? err.message : String(err);
log.warning(`» failed to unload model ${model}: ${msg}`); log.warning(`» failed to unload model ${model}: ${msg}`);
} }
} }
async function runOllamaLoop(ctx: AgentRunContext): Promise<AgentResult> { async function runOllamaLoop(ctx: AgentRunContext): Promise<AgentResult> {
const ollamaHost = process.env.OLLAMA_HOST ?? ""; const ollamaHost = process.env.OLLAMA_HOST ?? "";
if (!ollamaHost) { if (!ollamaHost) {
const errorMsg = const errorMsg =
"OLLAMA_HOST environment variable is not set. Please set it to the URL of your Ollama instance."; "OLLAMA_HOST environment variable is not set. Please set it to the URL of your Ollama instance.";
log.error(errorMsg); log.error(errorMsg);
return { success: false, error: errorMsg }; return { success: false, error: errorMsg };
} }
const model = ctx.model ?? process.env.OLLAMA_MODEL ?? DEFAULT_MODEL; const model = ctx.model ?? process.env.OLLAMA_MODEL ?? DEFAULT_MODEL;
const numCtx = ctx.payload.contextWindow ?? 262144; const numCtx = ctx.payload.contextWindow ?? 262144;
log.info(`» connecting to Ollama at ${ollamaHost}, model ${model}`); log.info(`» connecting to Ollama at ${ollamaHost}, model ${model}`);
const ollama = new Ollama({ host: ollamaHost }); const ollama = new Ollama({ host: ollamaHost });
const mcpClient = await buildMcpClient(ctx.mcpServerUrl); const mcpClient = await buildMcpClient(ctx.mcpServerUrl);
log.info("» fetching MCP tool list..."); log.info("» fetching MCP tool list...");
const tools = await getOllamaTools(mcpClient); const tools = await getOllamaTools(mcpClient);
log.info(`» ${tools.length} tools available`); log.info(`» ${tools.length} tools available`);
let messages: Message[] = [ let messages: Message[] = [
{
role: "user",
content: ctx.instructions.full,
},
];
// Tools that signal the agent has produced its final output.
const OUTPUT_TOOLS = new Set([
"create_pull_request_review",
"report_progress",
"set_output",
]);
let iterations = 0;
let pendingModeNudge = false;
let calledOutputTool = false;
let addedContinueNudge = false;
while (iterations < MAX_ITERATIONS) {
iterations++;
log.info(`» Ollama turn ${iterations}/${MAX_ITERATIONS}...`);
// Non-streaming with a heartbeat timer so the activity monitor stays alive
// during long prefill. Streaming was tried but Ollama only emits one tool
// call per chunk — batched tool calls collapse to one-per-turn, turning a
// 7-turn run into 26 turns. The heartbeat fires every 60s to prevent the
// 300s activity timeout from triggering during large-context prefill.
let response: Awaited<ReturnType<typeof ollama.chat>>;
const turnStart = Date.now();
const heartbeat = setInterval(() => {
log.info(`» still waiting for model... (${Math.round((Date.now() - turnStart) / 1000)}s)`);
}, 60_000);
try {
response = await retry(
() => ollama.chat({
model,
messages,
tools,
keep_alive: -1,
think: false,
options: { num_ctx: numCtx, temperature: 0.1 },
}),
{ {
delaysMs: [3_000, 8_000], role: "user",
shouldRetry: (err) => { content: ctx.instructions.full,
const msg = err instanceof Error ? err.message : String(err);
return /unexpected EOF|XML syntax error|ECONNRESET|ETIMEDOUT|fetch failed/i.test(msg);
},
label: `Ollama turn ${iterations}`,
}, },
); ];
} catch (err) {
clearInterval(heartbeat);
await unloadModel(ollama, model);
const lastError = err instanceof Error ? err.message : String(err);
log.error(`Ollama error: ${lastError}`);
return { success: false, error: `Ollama request failed: ${lastError}` };
}
clearInterval(heartbeat);
const promptTokens = response.prompt_eval_count; // Tools that signal the agent has produced its final output.
const evalTokens = response.eval_count; const OUTPUT_TOOLS = new Set([
const assistantMessage = response.message; "create_pull_request_review",
"report_progress",
"set_output",
]);
if (promptTokens !== undefined) { let iterations = 0;
const total = promptTokens + (evalTokens ?? 0); let pendingModeNudge = false;
const pct = Math.round((total / numCtx) * 100); let calledOutputTool = false;
log.info( let addedContinueNudge = false;
` context: ${promptTokens} prompt + ${evalTokens ?? 0} eval = ${total} tokens (${pct}% of ${numCtx} limit)`,
);
if (promptTokens > numCtx * 0.77) {
messages = pruneToolMessages(messages);
}
}
messages.push(assistantMessage); while (iterations < MAX_ITERATIONS) {
iterations++;
log.info(`» Ollama turn ${iterations}/${MAX_ITERATIONS}...`);
const toolCalls: ToolCall[] | undefined = assistantMessage.tool_calls; // Non-streaming with a heartbeat timer so the activity monitor stays alive
// during long prefill. Streaming was tried but Ollama only emits one tool
// call per chunk — batched tool calls collapse to one-per-turn, turning a
// 7-turn run into 26 turns. The heartbeat fires every 60s to prevent the
// 300s activity timeout from triggering during large-context prefill.
let response: Awaited<ReturnType<typeof ollama.chat>>;
const turnStart = Date.now();
const heartbeat = setInterval(() => {
log.info(`» still waiting for model... (${Math.round((Date.now() - turnStart) / 1000)}s)`);
}, 60_000);
try {
response = await retry(
() => ollama.chat({
model,
messages,
tools,
keep_alive: -1,
think: false,
options: { num_ctx: numCtx, temperature: 0.1 },
}),
{
delaysMs: [3_000, 8_000, 15_000, 30_000, 600_000], // up to 6 attempts over ~16 minutes
shouldRetry: (err) => {
const msg = err instanceof Error ? err.message : String(err);
return /unexpected EOF|XML syntax error|ECONNRESET|ETIMEDOUT|fetch failed/i.test(msg);
},
label: `Ollama turn ${iterations}`,
},
);
} catch (err) {
clearInterval(heartbeat);
await unloadModel(ollama, model);
const lastError = err instanceof Error ? err.message : String(err);
log.error(`Ollama error: ${lastError}`);
return { success: false, error: `Ollama request failed: ${lastError}` };
}
clearInterval(heartbeat);
if (!toolCalls || toolCalls.length === 0) { const promptTokens = response.prompt_eval_count;
log.debug(` model text: ${assistantMessage.content?.slice(0, 500)}`); const evalTokens = response.eval_count;
const assistantMessage = response.message;
// If the model stopped before ever calling an output tool and we haven't if (promptTokens !== undefined) {
// nudged yet, give it one more push to continue the workflow — regardless const total = promptTokens + (evalTokens ?? 0);
// of whether the mode nudge is still pending (the model may have stopped const pct = Math.round((total / numCtx) * 100);
// right after select_mode before acting on the guidance). log.info(
if (!calledOutputTool && !addedContinueNudge) { ` context: ${promptTokens} prompt + ${evalTokens ?? 0} eval = ${total} tokens (${pct}% of ${numCtx} limit)`,
log.info( );
"» model stopped before completing task — nudging to continue", if (promptTokens > numCtx * 0.77) {
messages = pruneToolMessages(messages);
}
}
messages.push(assistantMessage);
const toolCalls: ToolCall[] | undefined = assistantMessage.tool_calls;
if (!toolCalls || toolCalls.length === 0) {
log.debug(` model text: ${assistantMessage.content?.slice(0, 500)}`);
// If the model stopped before ever calling an output tool and we haven't
// nudged yet, give it one more push to continue the workflow — regardless
// of whether the mode nudge is still pending (the model may have stopped
// right after select_mode before acting on the guidance).
if (!calledOutputTool && !addedContinueNudge) {
log.info(
"» model stopped before completing task — nudging to continue",
);
addedContinueNudge = true;
messages.push({
role: "user",
content:
"Your task is not complete yet. Continue executing the workflow — " +
"call the next required tool to finish. " +
"Do not stop until you have submitted a review (create_pull_request_review) " +
"or called report_progress with a final summary.",
});
continue;
}
await unloadModel(ollama, model);
if (pendingModeNudge) {
log.info("» agent finished after mode nudge (no tool calls)");
} else {
log.info("» agent finished (no tool calls)");
}
return {
success: true,
output: assistantMessage.content || undefined,
};
}
pendingModeNudge = false;
const calledSelectMode = toolCalls.some(
(tc) => tc.function.name === "select_mode",
); );
addedContinueNudge = true;
messages.push({
role: "user",
content:
"Your task is not complete yet. Continue executing the workflow — " +
"call the next required tool to finish. " +
"Do not stop until you have submitted a review (create_pull_request_review) " +
"or called report_progress with a final summary.",
});
continue;
}
await unloadModel(ollama, model); for (const toolCall of toolCalls) {
const toolName = toolCall.function.name;
const toolArgs = toolCall.function.arguments;
if (pendingModeNudge) { log.info(`» calling tool: ${toolName}`);
log.info("» agent finished after mode nudge (no tool calls)"); log.debug(` args: ${JSON.stringify(toolArgs)}`);
} else {
log.info("» agent finished (no tool calls)");
}
return { if (OUTPUT_TOOLS.has(toolName)) {
success: true, calledOutputTool = true;
output: assistantMessage.content || undefined, }
};
if (ctx.onToolUse) {
ctx.onToolUse({ toolName, input: toolArgs });
}
const result = await callMcpTool(
mcpClient,
toolName,
toolArgs as Record<string, unknown>,
);
log.debug(` result: ${result.slice(0, 200)}`);
messages.push({
role: "tool",
content: result,
});
}
// After the FIRST select_mode call, nudge the model to act on the guidance.
// Only nudge once — repeated nudging causes a loop where the model keeps
// re-calling select_mode instead of executing the workflow.
if (calledSelectMode && !pendingModeNudge) {
pendingModeNudge = true;
// Parse the selected mode name from the tool result so we can give a
// more specific first-step instruction.
let selectedMode = "";
try {
const lastToolMsg = messages[messages.length - 1];
const parsed = JSON.parse(
typeof lastToolMsg.content === "string" ? lastToolMsg.content : "",
);
if (typeof parsed?.modeName === "string")
selectedMode = parsed.modeName;
} catch {
// best-effort
}
const firstStep =
selectedMode === "Review" || selectedMode === "IncrementalReview"
? "Your first tool call must be checkout_pr with the PR number from the event context."
: "Call the first tool required by the workflow now.";
messages.push({
role: "user",
content:
`Good. You have selected ${selectedMode || "a"} mode and received the workflow. ` +
"Do NOT call select_mode again. " +
`${firstStep} ` +
"Execute the complete workflow step by step until you call create_pull_request_review or report_progress.",
});
}
} }
pendingModeNudge = false; await unloadModel(ollama, model);
const calledSelectMode = toolCalls.some( log.warning(`» agent hit max iterations (${MAX_ITERATIONS})`);
(tc) => tc.function.name === "select_mode",
);
for (const toolCall of toolCalls) { return {
const toolName = toolCall.function.name; success: false,
const toolArgs = toolCall.function.arguments; error: `Agent exceeded maximum iterations (${MAX_ITERATIONS})`,
};
log.info(`» calling tool: ${toolName}`);
log.debug(` args: ${JSON.stringify(toolArgs)}`);
if (OUTPUT_TOOLS.has(toolName)) {
calledOutputTool = true;
}
if (ctx.onToolUse) {
ctx.onToolUse({ toolName, input: toolArgs });
}
const result = await callMcpTool(
mcpClient,
toolName,
toolArgs as Record<string, unknown>,
);
log.debug(` result: ${result.slice(0, 200)}`);
messages.push({
role: "tool",
content: result,
});
}
// After the FIRST select_mode call, nudge the model to act on the guidance.
// Only nudge once — repeated nudging causes a loop where the model keeps
// re-calling select_mode instead of executing the workflow.
if (calledSelectMode && !pendingModeNudge) {
pendingModeNudge = true;
// Parse the selected mode name from the tool result so we can give a
// more specific first-step instruction.
let selectedMode = "";
try {
const lastToolMsg = messages[messages.length - 1];
const parsed = JSON.parse(
typeof lastToolMsg.content === "string" ? lastToolMsg.content : "",
);
if (typeof parsed?.modeName === "string")
selectedMode = parsed.modeName;
} catch {
// best-effort
}
const firstStep =
selectedMode === "Review" || selectedMode === "IncrementalReview"
? "Your first tool call must be checkout_pr with the PR number from the event context."
: "Call the first tool required by the workflow now.";
messages.push({
role: "user",
content:
`Good. You have selected ${selectedMode || "a"} mode and received the workflow. ` +
"Do NOT call select_mode again. " +
`${firstStep} ` +
"Execute the complete workflow step by step until you call create_pull_request_review or report_progress.",
});
}
}
await unloadModel(ollama, model);
log.warning(`» agent hit max iterations (${MAX_ITERATIONS})`);
return {
success: false,
error: `Agent exceeded maximum iterations (${MAX_ITERATIONS})`,
};
} }
export const ollamaAgent = agent({ export const ollamaAgent = agent({
name: "ollama", name: "ollama",
run: async (ctx: AgentRunContext): Promise<AgentResult> => { run: async (ctx: AgentRunContext): Promise<AgentResult> => {
return runOllamaLoop(ctx); return runOllamaLoop(ctx);
}, },
}); });