speckl/reference-agent/src/orchestrator.ts

/**
 * Orchestrator — Client entry point for the Speckl agent harness.
 * 
 * Accepts Specks, decomposes them into verification and codegen tasks,
 * assigns tasks to Team Leads, and aggregates results with full provenance.
 */

import type {
  AgentId,
  SpeckDefinition,
  Task,
  TaskMessage,
  TaskResult,
  TaskType,
  ProvenanceEdge,
} from "./types.js";
import { TeamLead, VerificationTeamLead, CodegenTeamLead } from "./team-lead.js";
import { TillDoneTracker } from "./till-done.js";
import { ProvenanceGraphImpl } from "./provenance.js";

export class Orchestrator {
  readonly id: AgentId = { type: "orchestrator", name: "speckl-orchestrator" };
  private teamLeads: TeamLead[] = [];
  protected provenance: ProvenanceGraphImpl;
  protected tillDone: TillDoneTracker;
  private taskCounter = 0;

  constructor() {
    this.provenance = new ProvenanceGraphImpl();
    this.tillDone = new TillDoneTracker();

    // Register default teams
    this.registerTeamLead(new VerificationTeamLead());
    this.registerTeamLead(new CodegenTeamLead());
  }

  /** Register a Team Lead. */
  registerTeamLead(lead: TeamLead): void {
    this.teamLeads.push(lead);
  }

  /**
   * Process a Speck through the full pipeline:
   * 1. Verify constraints (Z3)
   * 2. Generate code (TypeScript/Rust)
   * 3. Return aggregated results with provenance graph
   */
  async processSpeck(speck: SpeckDefinition): Promise<OrchestrationResult> {
    const startTime = Date.now();
    const results: Map<string, TaskResult> = new Map();

    // Record speck in provenance graph
    this.provenance.addNode({
      id: `speck:${speck.name}`,
      type: "speck",
      label: speck.name,
      metadata: {
        constraintCount: speck.constraints.length,
        verifyCount: speck.verify.length,
        importedAt: new Date().toISOString(),
      },
    });

    // Phase 1: Verification
    const verifyTask = this.createTask(speck.name, "verify_speck", `Verify Speck ${speck.name} constraints`, 1);
    this.tillDone.register(verifyTask.id, `Z3 verifies ${speck.name}`, TillDoneTracker.predicates.z3Verified);

    const verifyLead = this.findTeamForTask(verifyTask);
    if (!verifyLead) {
      return this.fail(speck.name, "No verification team available", startTime);
    }

    const verifyResult = await verifyLead.handleAssignment(
      this.createAssignment(verifyTask, speck),
      speck
    );
    results.set(verifyTask.id, verifyResult);
    this.mergeProvenance(verifyResult.provenance);

    if (verifyResult.status !== "success") {
      return this.fail(speck.name, "Verification failed", startTime, results);
    }

    // Phase 2: Code generation
    const codegenTask = this.createTask(speck.name, "generate_code", `Generate code for ${speck.name}`, 2, verifyTask.id);
    this.tillDone.register(codegenTask.id, `Code generated for ${speck.name}`, TillDoneTracker.predicates.codeGenerated);

    const codegenLead = this.findTeamForTask(codegenTask);
    if (!codegenLead) {
      return this.fail(speck.name, "No codegen team available", startTime, results);
    }

    const codegenResult = await codegenLead.handleAssignment(
      this.createAssignment(codegenTask, speck),
      speck
    );
    results.set(codegenTask.id, codegenResult);
    this.mergeProvenance(codegenResult.provenance);

    const allDone = this.tillDone.completed().length === 2;

    return {
      speckName: speck.name,
      status: allDone ? "success" : "partial",
      results: Object.fromEntries(results),
      provenance: this.provenance.toJSON(),
      durationMs: Date.now() - startTime,
    };
  }

  /** Create a new task. */
  private createTask(
    speckName: string,
    taskType: TaskType,
    description: string,
    priority: number,
    parentTaskId?: string
  ): Task {
    return {
      id: `task-${++this.taskCounter}`,
      speckName,
      taskType,
      description,
      priority,
      parentTaskId,
      createdAt: new Date().toISOString(),
    };
  }

  /** Create an assignment message for a Team Lead. */
  private createAssignment(task: Task, speck: SpeckDefinition): TaskMessage {
    return {
      id: `msg-${task.id}`,
      type: "task_assign",
      from: this.id,
      to: { type: "team_lead", name: "auto" },
      payload: { task, speck },
      timestamp: new Date().toISOString(),
      provenance: [],
    };
  }

  /** Find a Team Lead that can handle the given task. */
  private findTeamForTask(task: Task): TeamLead | undefined {
    return this.teamLeads.find((lead) => lead.canHandle(task));
  }

  /** Merge provenance edges from a result. */
  private mergeProvenance(edges: ProvenanceEdge[]): void {
    for (const edge of edges) {
      this.provenance.addEdge(edge);
    }
  }

  /** Build a failure result. */
  private fail(
    speckName: string,
    reason: string,
    startTime: number,
    partialResults?: Map<string, TaskResult>
  ): OrchestrationResult {
    return {
      speckName,
      status: "failure",
      results: partialResults ? Object.fromEntries(partialResults) : {},
      provenance: this.provenance.toJSON(),
      durationMs: Date.now() - startTime,
      error: reason,
    };
  }

  /** Get the full provenance graph. */
  getProvenance(): object {
    return this.provenance.toJSON();
  }
}

export interface OrchestrationResult {
  speckName: string;
  status: "success" | "partial" | "failure";
  results: Record<string, TaskResult>;
  provenance: object;
  durationMs: number;
  error?: string;
}