panopticon/src/session.ts

/**
 * Shared session creation utilities for orchestrator, workers, and synthesizer.
 */
import {
  createAgentSession,
  DefaultResourceLoader,
  SessionManager,
  SettingsManager,
  AuthStorage,
  ModelRegistry,
  readOnlyTools,
  type AgentSession,
  type AgentSessionEvent,
} from "@mariozechner/pi-coding-agent";
import { getModel } from "@mariozechner/pi-ai";
import type { ThinkingLevel } from "@mariozechner/pi-agent-core";
import type { Config, SessionMetrics } from "./types.js";

export interface SessionOptions {
  role: "orchestrator" | "worker" | "synthesizer";
  projectPath: string;
  config: Config;
  systemPrompt: string;
}

/**
 * Resolve a model spec like "anthropic/claude-sonnet-4-5" into a Model object.
 */
function resolveModel(spec: string, modelRegistry: ModelRegistry) {
  const [provider, modelId] = spec.split("/");
  // Try ModelRegistry first, fall back to getModel
  try {
    return getModel(provider as any, modelId as any);
  } catch {
    throw new Error(`Cannot resolve model: ${spec}`);
  }
}

/**
 * Create a pi AgentSession for a given role.
 */
export async function createSession(options: SessionOptions): Promise<{
  session: AgentSession;
  metrics: SessionMetrics;
}> {
  const { role, projectPath, config, systemPrompt } = options;

  const authStorage = AuthStorage.create();
  const modelRegistry = new ModelRegistry(authStorage);

  const modelSpec = config.models[role];
  const model = resolveModel(modelSpec, modelRegistry);
  const thinkingLevel: ThinkingLevel = config.thinkingLevels[role];

  const loader = new DefaultResourceLoader({
    cwd: projectPath,
    systemPrompt,
    noSkills: true,
    noPromptTemplates: true,
    noExtensions: true,
    noThemes: true,
  });
  await loader.reload();

  const settingsManager = SettingsManager.inMemory({
    retry: { enabled: true, maxRetries: 2 },
  });

  // Workers get read-only tools; orchestrator and synthesizer get no tools
  const tools = role === "worker" ? readOnlyTools : [];

  const { session } = await createAgentSession({
    cwd: projectPath,
    model,
    thinkingLevel,
    tools,
    resourceLoader: loader,
    sessionManager: SessionManager.inMemory(),
    settingsManager,
    authStorage,
    modelRegistry,
  });

  // Track metrics
  const metrics = trackSession(session);

  return { session, metrics };
}

/**
 * Subscribe to session events and collect metrics.
 */
function trackSession(session: AgentSession): SessionMetrics {
  const metrics: SessionMetrics = {
    tokensIn: 0,
    tokensOut: 0,
    toolCalls: 0,
    errors: [],
    durationMs: 0,
  };

  session.subscribe((event: AgentSessionEvent) => {
    if (event.type === "message_end") {
      const msg = event.message as any;
      if (msg.role === "assistant" && msg.usage) {
        metrics.tokensIn += msg.usage.input ?? 0;
        metrics.tokensOut += msg.usage.output ?? 0;
      }
    }
    if (event.type === "tool_execution_end") {
      metrics.toolCalls++;
      if (event.isError) {
        metrics.errors.push(`${event.toolName}: ${JSON.stringify(event.result).slice(0, 200)}`);
      }
    }
  });

  return metrics;
}

/**
 * Extract the final text response from a session after prompt() completes.
 */
export function extractFinalResponse(session: AgentSession): string {
  const messages = session.messages;
  for (let i = messages.length - 1; i >= 0; i--) {
    const msg = messages[i] as any;
    if (msg.role === "assistant" && msg.content) {
      const textParts = msg.content.filter((c: any) => c.type === "text");
      return textParts.map((c: any) => c.text).join("\n");
    }
  }
  return "";
}

/**
 * Run a prompt with a timeout via AbortController.
 */
export async function promptWithTimeout(
  session: AgentSession,
  prompt: string,
  timeoutMs: number
): Promise<string> {
  const controller = new AbortController();
  const timeout = setTimeout(() => controller.abort(), timeoutMs);

  try {
    await session.prompt(prompt);
    return extractFinalResponse(session);
  } finally {
    clearTimeout(timeout);
  }
}