Files
ai_ops/src/examples/session-context.ts

106 lines
2.8 KiB
TypeScript

import { getConfig, type AppConfig } from "../config.js";
import type { AgentSession } from "../agents/manager.js";
import type { ProvisionedResources } from "../agents/provisioning.js";
import {
getAgentManager,
getResourceProvisioningOrchestrator,
} from "../agents/runtime.js";
import {
getDefaultMcpRegistry,
loadMcpConfigFromEnv,
type LoadedMcpConfig,
type McpRegistry,
} from "../mcp.js";
export type SessionProvider = "codex" | "claude";
export type SessionContext = {
provider: SessionProvider;
sessionId: string;
mcp: LoadedMcpConfig;
promptWithContext: string;
runtimeInjection: Awaited<ReturnType<ProvisionedResources["buildRuntimeInjection"]>>;
runInSession: <T>(run: () => Promise<T>) => Promise<T>;
close: () => Promise<void>;
};
export async function createSessionContext(
provider: SessionProvider,
input: {
prompt: string;
config?: Readonly<AppConfig>;
mcpRegistry?: McpRegistry;
},
): Promise<SessionContext> {
const config = input.config ?? getConfig();
const mcpRegistry = input.mcpRegistry ?? getDefaultMcpRegistry();
const agentManager = getAgentManager(config);
const agentSession = agentManager.createSession();
const resourceProvisioning = getResourceProvisioningOrchestrator(config);
let provisionedResources: ProvisionedResources | undefined;
let closed = false;
const close = async (): Promise<void> => {
if (closed) {
return;
}
closed = true;
if (provisionedResources) {
await provisionedResources.release();
}
agentSession.close();
};
try {
provisionedResources = await resourceProvisioning.provisionSession({
sessionId: agentSession.id,
resources: [{ kind: "git-worktree" }, { kind: "port-range" }],
});
const runtimeInjection = await provisionedResources.buildRuntimeInjection({
discoveryFileRelativePath: config.discovery.fileRelativePath,
baseEnv: process.env,
});
const promptWithContext = provisionedResources.composePrompt(input.prompt, [
`Discovery file: ${runtimeInjection.discoveryFilePath}`,
"Resource env vars are pre-injected (AGENT_WORKTREE_PATH, AGENT_PORT_RANGE_START, AGENT_PORT_RANGE_END, AGENT_PORT_PRIMARY).",
]);
const mcp = loadMcpConfigFromEnv(
{
providerHint: provider,
prompt: input.prompt,
},
{
config,
registry: mcpRegistry,
},
);
return {
provider,
sessionId: agentSession.id,
mcp,
promptWithContext,
runtimeInjection,
runInSession: <T>(run: () => Promise<T>) =>
runWithAgentSession(agentSession, run),
close,
};
} catch (error) {
await close();
throw error;
}
}
async function runWithAgentSession<T>(agentSession: AgentSession, run: () => Promise<T>): Promise<T> {
return agentSession.runAgent({
depth: 0,
run,
});
}