354 lines
14 KiB
TypeScript
354 lines
14 KiB
TypeScript
// The agent module's adapters: the procedure handlers, the tools gateway over the registry client,
|
|
// the run job that drives the engine, and the subscriber that enqueues it.
|
|
|
|
import {
|
|
AGENT_MODEL,
|
|
AGENT_TOOLS,
|
|
type AgentStep,
|
|
dataAgentPersistence,
|
|
dataRunStore,
|
|
executeRun,
|
|
expireRuns,
|
|
RunAuthorityDenied,
|
|
} from "@conusai-kit/agent"
|
|
import { runEventsResponse } from "@conusai-kit/agent/agui"
|
|
import type { Container } from "@conusai-kit/core"
|
|
import {
|
|
type EventContext,
|
|
implement,
|
|
internal,
|
|
isSystemActor,
|
|
JobClaimLost,
|
|
type JobDef,
|
|
subscribe,
|
|
} from "@conusai-kit/core"
|
|
import { DATA_SOURCE, JOB_QUEUE } from "@conusai-kit/data"
|
|
import { assertJobTransactionFence } from "@conusai-kit/jobs-pgboss"
|
|
import { TENANCY } from "@conusai-kit/tenancy"
|
|
import { createRegistryClient, type RegistryClient, ToolDeniedError } from "@conusai-sdk/ai"
|
|
import { type ConusaiTextAdapter, conusaiText } from "@conusai-sdk/ai/tanstack"
|
|
import { z } from "zod"
|
|
import {
|
|
approvalList,
|
|
approvalReview,
|
|
RUN_EXPIRE_JOB,
|
|
RUN_JOB,
|
|
runCancel,
|
|
runEvents,
|
|
runGetMany,
|
|
runList,
|
|
runResumeRequested,
|
|
runStart,
|
|
runStatus,
|
|
} from "../contract/index.js"
|
|
import { AGENT_SCOPE, type AgentActionScope, type AgentDeps, type AgentScope } from "../core/ports.js"
|
|
import {
|
|
cancelRun,
|
|
getManyRuns,
|
|
getRun,
|
|
listApprovals,
|
|
listRuns,
|
|
openRunEvents,
|
|
reviewApproval,
|
|
startRun,
|
|
} from "../core/runs.js"
|
|
|
|
export function agentDeps(container: Container): AgentDeps {
|
|
return { data: container.resolve(DATA_SOURCE) }
|
|
}
|
|
|
|
export const runStartImpl = implement(runStart, ({ input, ctx }) =>
|
|
startRun(agentDeps(ctx.container), ctx.scope, {
|
|
goal: input.goal,
|
|
plan: input.plan,
|
|
...(input.budget === undefined ? {} : { budget: input.budget }),
|
|
}),
|
|
)
|
|
|
|
export const runStatusImpl = implement(runStatus, ({ input, ctx }) =>
|
|
getRun(agentDeps(ctx.container), ctx.scope, input.runId),
|
|
)
|
|
|
|
/**
|
|
* The run's events as SSE: the request is authorized and the run found before the first byte, and
|
|
* the body then tails the durable record until the run's terminal event or the end of the stream
|
|
* window, when the client reconnects (and is authorized again) with its last offset. Each poll
|
|
* also rechecks current membership and permission without the request memo.
|
|
*/
|
|
export const runEventsImpl = implement(runEvents, async ({ input, ctx }) => {
|
|
const opened = await openRunEvents(agentDeps(ctx.container), ctx.scope, { runId: input.runId, after: input.after })
|
|
return runEventsResponse({
|
|
runId: input.runId,
|
|
after: opened.after,
|
|
load: async () => {
|
|
// Request authorization is memoised. A stream must observe a removal or demotion before
|
|
// reading another durable snapshot, rather than retaining that memo until reconnect.
|
|
const data = ctx.container.resolve(DATA_SOURCE)
|
|
const tenancy = ctx.container.resolve(TENANCY)
|
|
const allowed = await data.unitOfWork.run(ctx.scope, (tx) =>
|
|
tenancy.authorizeIn(tx, ctx.scope, runEvents.permission, "effect", runEvents.name),
|
|
)
|
|
return allowed ? opened.load() : null
|
|
},
|
|
...(ctx.signal === undefined ? {} : { signal: ctx.signal }),
|
|
})
|
|
})
|
|
|
|
export const runCancelImpl = implement(runCancel, ({ input, ctx }) =>
|
|
cancelRun(agentDeps(ctx.container), ctx.scope, input.runId),
|
|
)
|
|
|
|
export const runListImpl = implement(runList, ({ input, ctx }) =>
|
|
listRuns(agentDeps(ctx.container), ctx.scope, {
|
|
...(input.cursor === undefined ? {} : { cursor: input.cursor }),
|
|
...(input.limit === undefined ? {} : { limit: input.limit }),
|
|
}),
|
|
)
|
|
|
|
export const approvalListImpl = implement(approvalList, ({ ctx }) => listApprovals(agentDeps(ctx.container), ctx.scope))
|
|
|
|
export const approvalReviewImpl = implement(approvalReview, ({ input, ctx }) =>
|
|
reviewApproval(agentDeps(ctx.container), ctx.scope, input),
|
|
)
|
|
|
|
export const runGetManyImpl = implement(runGetMany, async ({ input, ctx }) => ({
|
|
items: await getManyRuns(agentDeps(ctx.container), ctx.scope, input.ids),
|
|
}))
|
|
|
|
export const AGENT_PROCEDURE_IMPLS = [
|
|
runStartImpl,
|
|
runStatusImpl,
|
|
runEventsImpl,
|
|
runCancelImpl,
|
|
runListImpl,
|
|
runGetManyImpl,
|
|
approvalListImpl,
|
|
approvalReviewImpl,
|
|
]
|
|
|
|
// --- The model -------------------------------------------------------------
|
|
|
|
/**
|
|
* The run loop's model: the AI gateway's `model` as a TanStack AI text adapter, with the app's
|
|
* deployment credentials. The worker binds it as `AGENT_MODEL`; the gateway meters every call and
|
|
* applies the organization's spend cap, and no provider SDK or key is in the app.
|
|
*/
|
|
export function gatewayAgentModel(options: {
|
|
apiUrl: string
|
|
token: string
|
|
model: string
|
|
}): ConusaiTextAdapter<string> {
|
|
return conusaiText(options.model, { apiUrl: options.apiUrl, token: options.token })
|
|
}
|
|
|
|
// --- The tools gateway -----------------------------------------------------
|
|
|
|
/**
|
|
* A tool call no retry fixes by itself: the platform refused it, or the tool answered an error. The
|
|
* run job turns it into the step's typed failure. Its message goes on the run, so it names the tool
|
|
* and the refusal's reason and never the tool's own text, which may be a tenant's data.
|
|
*/
|
|
class ToolCallFailed extends Error {}
|
|
|
|
/**
|
|
* The tools gateway over the organization's registry, with the app's scoped credentials. Each
|
|
* step's call carries the step's idempotency key, so a resumed step re-calls the same persisted
|
|
* action instead of creating a second one. When a write comes back `approval_required`, the exact
|
|
* stored proposed arguments are parked for a person; this adapter never approves. A refusal and a
|
|
* tool's own error are `ToolCallFailed`; anything else that throws is transient, and leaves the run to
|
|
* the queue's retry.
|
|
*/
|
|
export function registryAgentTools(options: {
|
|
apiUrl: string
|
|
token: string
|
|
/** The organization's id the registry is addressed by. */
|
|
orgId: string
|
|
/** Fixed, or resolved per call: jobs of different tenants share the process binding. */
|
|
scope: AgentActionScope | (() => Promise<AgentActionScope>)
|
|
fetch?: typeof fetch
|
|
}): import("@conusai-kit/agent").AgentTools {
|
|
// Lazy: building the client validates the organization id, and a bare test container binds the
|
|
// tools without one until a call actually leaves the process.
|
|
let registry: RegistryClient | undefined
|
|
const client = (): RegistryClient => {
|
|
registry ??= createRegistryClient({
|
|
env: {
|
|
CONUSAI_API_URL: options.apiUrl,
|
|
CONUSAI_API_TOKEN: options.token,
|
|
CONUSAI_ORG_ID: options.orgId,
|
|
},
|
|
...(options.fetch === undefined ? {} : { fetch: options.fetch }),
|
|
})
|
|
return registry
|
|
}
|
|
const resolveScope =
|
|
typeof options.scope === "function" ? options.scope : () => Promise.resolve(options.scope as AgentActionScope)
|
|
return {
|
|
async call(name, args, callOptions) {
|
|
const result = await client()
|
|
.tools.call(name, args, { idempotencyKey: callOptions.idempotencyKey })
|
|
.catch((error: unknown) => {
|
|
throw error instanceof ToolDeniedError && !error.retryable ? new ToolCallFailed(error.message) : error
|
|
})
|
|
if (result.status === "completed") {
|
|
if (result.isError) throw new ToolCallFailed(`the tool ${name} answered an error`)
|
|
return {
|
|
status: "completed",
|
|
content: result.content,
|
|
// The registry call result does not carry metered usage yet (the metering slice owns
|
|
// that); steps and duration are tracked, and cost lands when the platform answers it.
|
|
usage: { costMicrocents: 0, steps: 1, tokens: 0 },
|
|
}
|
|
}
|
|
return {
|
|
status: "approval_required",
|
|
approval: {
|
|
actionId: result.request.actionId ?? null,
|
|
tool: name,
|
|
arguments: result.request.arguments,
|
|
reviewUrl: result.request.actionId
|
|
? client()
|
|
.tools.actions(await resolveScope())
|
|
.reviewUrl(result.request.actionId)
|
|
: null,
|
|
},
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
// --- The run job and its subscriber ----------------------------------------
|
|
|
|
const runPayload = z.object({ runId: z.string() })
|
|
type RunJobPayload = z.output<typeof runPayload>
|
|
|
|
export const runJob: JobDef<RunJobPayload> = {
|
|
payload: runPayload,
|
|
async handler(payload, ctx) {
|
|
const resolved = ctx.container.resolve(AGENT_SCOPE)
|
|
if (resolved.level !== "project")
|
|
throw internal(`an agent job runs in a project scope, not a ${resolved.level} one`)
|
|
const execution = ctx.execution
|
|
const actual = execution?.scope
|
|
if (
|
|
execution?.adapter !== "pg-boss" ||
|
|
execution.job !== RUN_JOB ||
|
|
!actual ||
|
|
actual.level !== "project" ||
|
|
!isSystemActor(actual.actor) ||
|
|
!isSystemActor(resolved.actor) ||
|
|
actual.actor.userId !== resolved.actor.userId ||
|
|
actual.orgId !== resolved.orgId ||
|
|
actual.workspaceId !== resolved.workspaceId ||
|
|
actual.projectId !== resolved.projectId
|
|
)
|
|
throw new JobClaimLost()
|
|
return execution.run(async () => {
|
|
const data = ctx.container.resolve(DATA_SOURCE)
|
|
await assertJobTransactionFence(() => data.unitOfWork.run(resolved, async () => {}))
|
|
await execution.assertOwned()
|
|
const store = dataRunStore(data, resolved)
|
|
const record = await store.get(payload.runId)
|
|
if (!record) throw internal("the run this job belongs to is gone")
|
|
if (execution.onBehalfOf !== record.actorId) throw new JobClaimLost()
|
|
const authorize = async (): Promise<void> => {
|
|
await execution.assertOwned()
|
|
// The worker actor is a system actor. Only the immutable initiating member can authorize
|
|
// this run, and every boundary reads current membership instead of request memo state.
|
|
const member = { ...resolved, actor: { userId: record.actorId } }
|
|
const tenancy = ctx.container.resolve(TENANCY)
|
|
const allowed = await data.unitOfWork.run(member, (tx) =>
|
|
tenancy.authorizeIn(tx, member, runStart.permission, "effect", runStart.name),
|
|
)
|
|
await execution.assertOwned()
|
|
if (!allowed) throw new RunAuthorityDenied()
|
|
}
|
|
const plan: AgentStep[] = record.plan.map((planned) => ({
|
|
key: planned.key,
|
|
label: planned.label,
|
|
...(planned.needsApproval === true
|
|
? { requiresApproval: { tool: planned.tool, arguments: planned.args } }
|
|
: {}),
|
|
invoke: async (stepCtx) => {
|
|
await authorize()
|
|
const result = await ctx.container
|
|
.resolve(AGENT_TOOLS)
|
|
.call(planned.tool, planned.args, {
|
|
signal: stepCtx.signal,
|
|
idempotencyKey: `${record.id}:${planned.key}`,
|
|
})
|
|
.catch((error: unknown) => {
|
|
if (error instanceof ToolCallFailed) return { status: "failed" as const, detail: error.message }
|
|
throw error
|
|
})
|
|
await execution.assertOwned()
|
|
if (result.status === "failed") {
|
|
return { ok: false as const, failure: "tool_error" as const, detail: result.detail }
|
|
}
|
|
if (result.status === "completed") {
|
|
return {
|
|
ok: true as const,
|
|
usage: { costMicrocents: result.usage.costMicrocents, tokens: result.usage.tokens },
|
|
artifacts: [{ kind: "tool_result", ref: planned.tool }],
|
|
result: result.content,
|
|
}
|
|
}
|
|
return { ok: false as const, awaitingApproval: result.approval }
|
|
},
|
|
}))
|
|
await executeRun(
|
|
{
|
|
store,
|
|
authorize,
|
|
assertOwned: () => execution.assertOwned(),
|
|
// Resolve provider bindings only after entry admission; a denied/terminal run never
|
|
// constructs a gateway client. Metering and spend admission remain in the gateway.
|
|
get adapter() {
|
|
return ctx.container.resolve(AGENT_MODEL)
|
|
},
|
|
persistence: dataAgentPersistence(data, resolved),
|
|
},
|
|
payload.runId,
|
|
plan,
|
|
{
|
|
...(ctx.finalAttempt === undefined ? {} : { finalAttempt: ctx.finalAttempt }),
|
|
signal: ctx.signal ? AbortSignal.any([ctx.signal, execution.signal]) : execution.signal,
|
|
},
|
|
)
|
|
})
|
|
},
|
|
}
|
|
|
|
async function onRunResumeRequested(
|
|
event: { name: string; payload: { workspaceId: string; projectId: string; runId: string } },
|
|
ctx: EventContext,
|
|
): Promise<void> {
|
|
const ambient = ctx.container.resolve(AGENT_SCOPE)
|
|
if (ambient.level !== "project") throw internal("an agent event resolves to a project scope")
|
|
const scope: AgentScope = { ...ambient, workspaceId: event.payload.workspaceId, projectId: event.payload.projectId }
|
|
await ctx.container
|
|
.resolve(JOB_QUEUE)
|
|
.enqueue(scope, RUN_JOB, { runId: event.payload.runId }, { key: `agent.run.${event.payload.runId}` })
|
|
}
|
|
|
|
const runExpirePayload = z.object({ retentionDays: z.number().int().min(1).max(3650) })
|
|
type RunExpirePayload = z.output<typeof runExpirePayload>
|
|
|
|
/**
|
|
* The retention sweep: deletes the finished runs of one project that are past the retention, with their step
|
|
* log and conversation (`expireRuns`). Runs in the project the schedule fanned it out to, so another
|
|
* organization's runs are not reached; a suspended organization gets no run at all (the fan-out skips it).
|
|
*/
|
|
export const runExpireJob: JobDef<RunExpirePayload> = {
|
|
payload: runExpirePayload,
|
|
async handler(payload, ctx) {
|
|
const resolved = ctx.container.resolve(AGENT_SCOPE)
|
|
if (resolved.level !== "project")
|
|
throw internal(`the run retention sweep runs in a project scope, not a ${resolved.level} one`)
|
|
await expireRuns(ctx.container.resolve(DATA_SOURCE), resolved, { retentionDays: payload.retentionDays })
|
|
},
|
|
}
|
|
|
|
export const AGENT_JOBS: Record<string, JobDef<unknown>> = { [RUN_JOB]: runJob, [RUN_EXPIRE_JOB]: runExpireJob }
|
|
export const AGENT_SUBSCRIPTIONS = [subscribe(runResumeRequested, onRunResumeRequested)]
|