Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions backend/src/knowledge/gateway/child-session-delegation.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -79,4 +79,30 @@ describe("child-session gateway control-plane delegation", () => {
expect(result.isError).toBe(true);
expect(result.content[0]?.text).toContain("not configured");
});

test("returns a structured control-plane refusal to the agent", async () => {
process.env.GATEWAY_DATABASE_URL = "postgres://restricted";
process.env.USEAGENT_API_ORIGIN = "http://127.0.0.1:3201";
globalThis.fetch = (async (_input, _init) => Response.json(
{
error: "fleet_queue_full",
message: "fleet queue is full for org org-a (limit 4)",
limit: 4,
},
{ status: 429 },
)) as typeof fetch;

const result = await executeChildSessionTool(claims, "child_session_create", {
idempotencyKey: "queue-full",
prompt: "do a thing",
});

expect(result.isError).toBe(true);
expect(result.content[0]?.text).toBe("fleet queue is full for org org-a (limit 4)");
expect(result.structuredContent).toEqual({
code: "fleet_queue_full",
status: 429,
limit: 4,
});
});
});
21 changes: 14 additions & 7 deletions backend/src/knowledge/gateway/child-session-internal-routes.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import { Hono } from "hono";
import {
CHILD_SESSION_TOOL_NAMES,
childSessionRefusalResult,
executeChildSessionToolLocal,
} from "./child-session-tools";
import { resolveToolRunIdentity } from "./run-authorization";
Expand Down Expand Up @@ -57,11 +58,17 @@ internalChildSessionRoutes.post("/", async (c) => {
return c.json({ error: "invalid_arguments" }, 400);
}

return c.json({
result: await executeChildSessionToolLocal(
current,
name,
args as Record<string, unknown>,
),
});
try {
return c.json({
result: await executeChildSessionToolLocal(
current,
name,
args as Record<string, unknown>,
),
});
} catch (error) {
const refusal = childSessionRefusalResult(error);
if (refusal) return c.json({ result: refusal });
throw error;
}
});
106 changes: 105 additions & 1 deletion backend/src/knowledge/gateway/child-session-tools.test.ts
Original file line number Diff line number Diff line change
@@ -1,17 +1,20 @@
import { afterEach, beforeEach, describe, expect, test } from "bun:test";
import { and, eq, isNotNull, isNull, like } from "drizzle-orm";
import "../../index";
import { commands, providerEvents, runs } from "../../db/schema";
import { commands, providerEvents, runs, spendAccounts } from "../../db/schema";
import { db } from "../../db/client";
import { createRun, setRunStatus } from "../../runs/repo";
import { recordProviderEvent } from "../../runs/provider-events";
import { createGatewayApp } from "../../gateway-app";
import { advertisedGatewayToolDescriptors, executeRegisteredGatewayTool } from "./operation-registry";
import {
CHILD_SESSION_TOOLS,
childSessionRefusalResult,
executeChildSessionTool,
} from "./child-session-tools";
import { mintToolToken, type ToolTokenClaims } from "./token";
import { RunAdmissionClosedError } from "../../commands/admission";
import { SandboxMinutesExceededError } from "../../runs/sandbox-minutes";

const touchedOrgs = new Set<string>();

Expand All @@ -25,6 +28,8 @@ const ENV_KEYS = [
"PRODUCT_CHILD_THREADS",
"GATEWAY_DATABASE_URL",
"USEAGENT_API_ORIGIN",
"FLEET_ORG_MAX_QUEUE_DEPTH",
"SPEND_ALLOWANCE_USD",
] as const;

let savedEnv: Partial<Record<(typeof ENV_KEYS)[number], string | undefined>> =
Expand Down Expand Up @@ -58,6 +63,7 @@ afterEach(async () => {
await db
.delete(runs)
.where(and(eq(runs.orgId, orgId), isNull(runs.parentRunId)));
await db.delete(spendAccounts).where(eq(spendAccounts.orgId, orgId));
}
touchedOrgs.clear();
for (const key of ENV_KEYS) {
Expand Down Expand Up @@ -112,6 +118,33 @@ function childId(
}

describe("child session gateway tools", () => {
test("refusal mapping preserves release and sandbox guidance", () => {
const release = childSessionRefusalResult(new RunAdmissionClosedError({
open: false,
operationId: "release-test",
actor: "test",
reason: "deploying",
changedAt: "2026-10-25T00:00:00.000Z",
}));
expect(release?.content[0]?.text).toBe(
"A release is being installed. Send your task again in a moment.",
);
expect(release?.structuredContent).toEqual({
code: "run_admission_closed",
retryable: true,
});

const sandbox = childSessionRefusalResult(
new SandboxMinutesExceededError(10, 10),
);
expect(sandbox?.content[0]?.text).toContain("10 of your 10 sandbox minutes");
expect(sandbox?.structuredContent).toEqual({
code: "sandbox_minutes_exceeded",
used: 10,
cap: 10,
});
});

test("batch advertisement is rollout honest", () => {
const names = () => advertisedGatewayToolDescriptors({ childSessions: true, slack: false }).map((tool) => tool.name);
expect(names()).not.toContain("child_session_create_many");
Expand Down Expand Up @@ -260,6 +293,77 @@ describe("child session gateway tools", () => {
expect(childCommands).toHaveLength(1);
});

test("create names the queue limit when the durable queue is full", async () => {
process.env.FLEET_ORG_MAX_QUEUE_DEPTH = "1";
const claims = await fixture();
const first = await executeChildSessionTool(claims, "child_session_create", {
idempotencyKey: "queue-first",
prompt: "Occupy the only queue slot.",
});
expect(first.isError).toBeUndefined();

const refused = await executeChildSessionTool(claims, "child_session_create", {
idempotencyKey: "queue-second",
prompt: "This child should be refused.",
});

expect(refused.isError).toBe(true);
expect(refused.content[0]?.text).toContain("fleet queue is full");
expect(refused.content[0]?.text).toContain("limit 1");
expect(refused.structuredContent).toMatchObject({
code: "fleet_queue_full",
retryable: true,
limit: 1,
});
});

test("create carries the spend allowance message", async () => {
process.env.SPEND_ALLOWANCE_USD = "5";
const claims = await fixture();
await db.insert(spendAccounts).values({
orgId: claims.orgId,
userId: claims.userId,
spentUsd: 5,
});

const refused = await executeChildSessionTool(claims, "child_session_create", {
idempotencyKey: "spend-refused",
prompt: "This child should be refused.",
});

expect(refused.isError).toBe(true);
expect(refused.content[0]?.text).toBe(
"You have spent $5.00 of your $5.00 allowance. New tasks are paused until it is raised.",
);
expect(refused.structuredContent).toMatchObject({
code: "spend_allowance_exceeded",
spent: 5,
allowance: 5,
});
});

test("create_many returns a structured queue refusal", async () => {
process.env.PRODUCT_CHILD_THREADS = "on";
process.env.FLEET_ORG_MAX_QUEUE_DEPTH = "1";
const claims = await fixture();

const refused = await executeChildSessionTool(claims, "child_session_create_many", {
idempotencyKey: "batch-queue-full",
children: [
{ title: "One", prompt: "First child." },
{ title: "Two", prompt: "Second child." },
],
});

expect(refused.isError).toBe(true);
expect(refused.content[0]?.text).toContain("fleet queue is full");
expect(refused.structuredContent).toMatchObject({
code: "fleet_queue_full",
retryable: true,
limit: 1,
});
});

test("list, events, and gather are bounded durable projections", async () => {
const claims = await fixture();
const first = await executeChildSessionTool(
Expand Down
81 changes: 65 additions & 16 deletions backend/src/knowledge/gateway/child-session-tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,10 @@ import type { GatewayToolListOptions } from "./operation-registry";
import { botsEnabled } from "../../bots/rollout";
import { productFlagsForToolList } from "./product-flags";
import { ENGINE_IDS, type EngineId } from "../../db/schema";
import { RunAdmissionClosedError } from "../../commands/admission";
import { FleetQueueLimitError } from "../../fleet/intake";
import { SandboxMinutesExceededError } from "../../runs/sandbox-minutes";
import { SpendAllowanceExceededError } from "../../runs/spend";

const MAX_TEXT_EVENT_LINES = 20;
const MAX_TEXT_PAYLOAD_CHARS = 320;
Expand Down Expand Up @@ -265,6 +269,37 @@ async function currentRun(claims: ToolTokenClaims) {
return run?.status === "running" ? run : null;
}

export function childSessionRefusalResult(error: unknown): ToolCallResult | null {
if (error instanceof FleetQueueLimitError) {
return errorResult(error.message, {
code: error.code,
retryable: true,
limit: error.limit,
});
}
if (error instanceof SpendAllowanceExceededError) {
return errorResult(error.message, {
code: error.code,
spent: error.spent,
allowance: error.allowance,
});
}
if (error instanceof SandboxMinutesExceededError) {
return errorResult(error.message, {
code: error.code,
used: error.used,
cap: error.cap,
});
}
if (error instanceof RunAdmissionClosedError) {
return errorResult(error.body.message, {
code: error.code,
retryable: true,
});
}
return null;
}

async function create(
claims: ToolTokenClaims,
args: Record<string, unknown>,
Expand All @@ -287,19 +322,26 @@ async function create(
);
}

const outcome = await createChildSession({
orgId: claims.orgId,
actorId: claims.userId || null,
parentRunId: run.id,
threadId: run.threadId,
prompt,
title: title || undefined,
engine: run.engine,
model: run.model,
repos: run.repos,
memoryScope: run.memoryScope,
idempotencyKey,
});
let outcome: Awaited<ReturnType<typeof createChildSession>>;
try {
outcome = await createChildSession({
orgId: claims.orgId,
actorId: claims.userId || null,
parentRunId: run.id,
threadId: run.threadId,
prompt,
title: title || undefined,
engine: run.engine,
model: run.model,
repos: run.repos,
memoryScope: run.memoryScope,
idempotencyKey,
});
} catch (error) {
const refusal = childSessionRefusalResult(error);
if (refusal) return refusal;
throw error;
}
if (outcome.status === "conflict") {
return errorResult(
"idempotencyKey was already used for different child session input.",
Expand Down Expand Up @@ -355,6 +397,8 @@ async function createMany(
outcome,
);
} catch (error) {
const refusal = childSessionRefusalResult(error);
if (refusal) return refusal;
return errorResult(error instanceof Error ? error.message : "child fan-out failed");
}
}
Expand Down Expand Up @@ -576,12 +620,17 @@ async function executeThroughPrimaryApi(
signal: AbortSignal.timeout(30_000),
});
const body = (await response.json().catch(() => null)) as
| { result?: ToolCallResult; error?: string }
| ({ result?: ToolCallResult; error?: string; message?: string } & Record<string, unknown>)
| null;
if (!response.ok || !body?.result) {
const { error, message, result: _result, ...details } = body ?? {};
return errorResult(
body?.error ?? `child-session control plane returned HTTP ${response.status}`,
{ status: response.status },
message ?? error ?? `child-session control plane returned HTTP ${response.status}`,
{
...(error ? { code: error } : {}),
status: response.status,
...details,
},
);
}
return body.result;
Expand Down
Loading