Luigit
repositories / will

will

owned by admin

src/agent/pi-session.ts

Raw
import { mkdir } from 'node:fs/promises'
import { join } from 'node:path'
import type { AgentSession, ToolDefinition } from '@earendil-works/pi-coding-agent'
import {
  createAgentSession,
  createCodingTools,
  createFindTool,
  createGrepTool,
  createLsTool,
  DefaultResourceLoader,
  ModelRuntime,
  SessionManager,
} from '@earendil-works/pi-coding-agent'
import type { AgentConfig } from './config.ts'
import type { SessionPort } from './engine.ts'

/**
 * The resident pi session adapter (spec AC15): one AgentSession per agent
 * process, resumed from the session file on the agent volume. Resource
 * discovery uses the immutable image while filesystem tools use the genome.
 */

export interface PiSessionOptions {
  config: AgentConfig
  systemPrompt: string
  customTools: ToolDefinition[]
  onEvent: (event: unknown) => void
}

export function createWorktreeTools(workDir: string, willEnv: Record<string, string>) {
  return [
    ...createCodingTools(workDir, {
      bash: {
        spawnHook: (context) => ({ ...context, env: { ...context.env, ...willEnv } }),
      },
    }),
    createGrepTool(workDir),
    createFindTool(workDir),
    createLsTool(workDir),
  ]
}

export class PiSession implements SessionPort {
  readonly #session: AgentSession

  private constructor(session: AgentSession) {
    this.#session = session
  }

  static async create(opts: PiSessionOptions): Promise<PiSession> {
    const cfg = opts.config
    if (!cfg.zai.apiKey)
      throw new Error('ZAI_API_KEY missing: cannot create the resident model session')

    const modelRuntime = await ModelRuntime.create({ modelsPath: null })
    await modelRuntime.setRuntimeApiKey('zai', cfg.zai.apiKey)
    const model =
      modelRuntime.getModel('zai', cfg.zai.analysisModel) ??
      modelRuntime.getModel('zai', 'glm-5.3-flash')
    if (!model) throw new Error(`zai model absent from built-in catalog (${cfg.zai.analysisModel})`)

    const loader = new DefaultResourceLoader({
      cwd: cfg.resourceDir,
      agentDir: join(cfg.agentDir, '.pi-agent'),
      systemPromptOverride: () => opts.systemPrompt,
      appendSystemPromptOverride: () => [],
    })
    await loader.reload()

    const willEnv: Record<string, string> = {
      WILL_WORK_DIR: cfg.workDir,
      WILL_RESOURCE_DIR: cfg.resourceDir,
      WILL_SERVICE_URL: cfg.serviceUrl,
      WILL_AGENT_DIR: cfg.agentDir,
      WILL_VCS_BRANCH: cfg.vcs.branch,
    }
    if (cfg.revision !== undefined) willEnv.WILL_REVISION = cfg.revision
    if (cfg.agentDigest !== undefined) willEnv.WILL_AGENT_DIGEST = cfg.agentDigest

    // Resource discovery stays in the immutable image; every filesystem tool
    // is explicitly rebound to the editable genome.
    const worktreeTools = createWorktreeTools(cfg.workDir, willEnv)

    await mkdir(cfg.agentDir, { recursive: true })
    const sessionManager = SessionManager.open(cfg.sessionFile)

    const { session } = await createAgentSession({
      cwd: cfg.resourceDir,
      agentDir: join(cfg.agentDir, '.pi-agent'),
      model,
      thinkingLevel: cfg.zai.thinkingLevel,
      modelRuntime,
      resourceLoader: loader,
      sessionManager,
      tools: [
        'read',
        'bash',
        'edit',
        'write',
        'grep',
        'find',
        'ls',
        ...opts.customTools.map((t) => t.name),
      ],
      customTools: [...worktreeTools, ...opts.customTools],
    })
    session.subscribe(opts.onEvent as never)
    return new PiSession(session)
  }

  async prompt(text: string): Promise<void> {
    await this.#session.prompt(text)
  }

  async abort(): Promise<void> {
    await this.#session.abort()
  }

  get sessionId(): string {
    return this.#session.sessionId
  }

  dispose(): void {
    this.#session.dispose()
  }
}

/**
 * Fake session for dev/e2e: records prompts, optionally auto-replies via a
 * callback. Enabled with WILL_SESSION_MODE=fake (never in production images).
 */
export function createFakeSession(
  onPrompt?: (text: string) => Promise<void>,
): SessionPort & { prompts: string[] } {
  const prompts: string[] = []
  return {
    prompts,
    async prompt(text: string) {
      prompts.push(text)
      await onPrompt?.(text)
    },
    async abort() {},
  }
}