@kontextmind/kxm/plugins/kxm/src/arbiter.ts

KXM local-first multi-agent orchestration and operator dashboard

import {
  DEFAULT_CONTEXT_BUDGET_TOKENS,
  MAX_CONTEXT_ITEMS,
  contextItemCharacters,
  estimateContextTokens,
  parseContextItem,
  parseContextRequest,
  provenanceSummaryOf,
  validateContextPacketContents,
  CONTEXT_SOURCE_TYPES,
  type ContextAuthority,
  type ContextItem,
  type ContextItemKind,
  type ContextPacket,
  type ContextRequest,
  type ContextSourceType,
} from "./context.ts";
import { ProtocolError } from "./protocol.ts";
import {
  compareCodeUnitIds,
  contextItemRelevanceText,
  relevanceTokens,
  roundRelevance,
  scoreRelevance,
} from "./relevance.ts";
import type { SkillLifecycle } from "./skills.ts";
import type { JournalCategory, WorkflowJournalEntry } from "./workflow.ts";
import type { MemoryRecord } from "./memory.ts";

/**
 * Role-aware context arbiter (v0.5, issue #34).
 *
 * The arbiter assembles context packets from the durable pool (journal
 * evidence, temporal state, knowledge, episodes, skills) for one role, task,
 * and token budget. Selection is deterministic: the same pool and request
 * always produce the same packet, so packets are testable, explainable, and
 * reproducible.
 */

export interface RoleContextPolicy {
  role: string;
  label: string;
  /** Item kinds this role receives, in priority order. */
  kinds: ContextItemKind[];
  /** Journal categories this role recalls from, in priority order. */
  journalCategories: JournalCategory[];
  /** Fixed default token budget for this role. */
  budgetTokens: number;
}

export const ROLE_POLICIES: readonly RoleContextPolicy[] = [
  {
    role: "repro",
    label: "Reproduction specialist",
    kinds: ["episode", "knowledge", "evidence"],
    journalCategories: ["error", "lesson", "observation", "contradiction"],
    budgetTokens: 8_000,
  },
  {
    role: "planner",
    label: "Planner",
    kinds: ["state", "knowledge", "evidence"],
    journalCategories: ["plan", "decision", "contradiction", "observation", "hypothesis", "experiment"],
    budgetTokens: 16_000,
  },
  {
    role: "critic",
    label: "Independent critic",
    kinds: ["knowledge", "evidence", "episode"],
    journalCategories: ["contradiction", "error", "lesson", "experiment"],
    budgetTokens: 12_000,
  },
  {
    role: "implementer",
    label: "Implementer",
    kinds: ["knowledge", "state", "skill", "episode", "evidence"],
    journalCategories: ["plan", "decision", "lesson", "state-change"],
    budgetTokens: 16_000,
  },
  {
    role: "verifier",
    label: "Verifier",
    kinds: ["evidence", "knowledge", "episode"],
    journalCategories: ["plan", "error", "lesson", "contradiction"],
    budgetTokens: 8_000,
  },
];

export function rolePolicy(role: string): RoleContextPolicy {
  return ROLE_POLICIES.find((policy) => policy.role === role) ?? {
    role,
    label: role,
    kinds: ["knowledge", "evidence"],
    journalCategories: ["lesson", "observation"],
    budgetTokens: DEFAULT_CONTEXT_BUDGET_TOKENS,
  };
}

const CONFIDENCE_RANK: Record<ContextItem["confidence"], number> = { verified: 3, probable: 2, uncertain: 1 };
const AUTHORITY_WEIGHT: Record<ContextItem["authority"], number> = { policy: 3, instruction: 2, evidence: 1, hypothesis: 0 };

export interface ArbiterOptions {
  /** Hub-owned working state (run/stage/attempt) delivered verbatim. */
  workingState?: Record<string, unknown>;
  /** Pool item IDs that represent open contradictions; they are routed to the
   * packet's contradictions section instead of their kind's section. */
  contradictionIds?: string[];
  /** Governed skill lifecycle (or a pre-verified snapshot of one) to read and
   * verify promoted skills by hash. */
  skillLifecycle?: Pick<SkillLifecycle, "list" | "verify">;
}

export interface ArbiterOutcome {
  packet: ContextPacket;
  /** Metadata-only telemetry: what was selected, from where, at what cost.
   * Never includes raw bodies. */
  audit: {
    request: { project: string; role: string; task: string; workflowRunId?: string; stageId?: string };
    selectedIds: string[];
    provenanceSummary: Record<string, number>;
    estimatedTokens: number;
    budgetTokens: number;
    candidateCount: number;
    excludedSuperseded: number;
    unresolvedGaps: string[];
    /** Numbers only: distinct task tokens, eligible candidates sharing a task
     * token, and each selected item's rounded BM25 score (index-aligned with
     * selectedIds). */
    relevance: { taskTokens: number; matchedCandidates: number; selected: number[] };
  };
}

/** Deterministically assemble a role-aware context packet. Superseded and
 * rejected records are excluded by default; cross-project content fails
 * closed; eligible candidates are ranked by contradiction, project-first,
 * task match, role kind priority, BM25 task relevance, confidence,
 * authority, recency (newest first), then id by code unit; the token budget
 * is filled first-fit on the serialized selection. */
export function arbitrate(
  requestInput: unknown,
  pool: ContextItem[],
  options: ArbiterOptions = {},
): ArbiterOutcome {
  const request = parseContextRequest(requestInput);
  const policy = rolePolicy(request.role);
  const budget = request.budgetTokens ?? policy.budgetTokens;
  const contradictions = new Set(options.contradictionIds ?? []);

  const candidates: ContextItem[] = [];
  let excludedSuperseded = 0;
  for (const candidate of pool) {
    if (candidate.project !== request.project && candidate.project !== "_shared") {
      // Defense in depth: the hub pre-filters by project, so a foreign item
      // in the pool is a bug — fail closed rather than silently filter.
      throw new ProtocolError(
        403,
        `context pool contains cross-project item ${candidate.id}`,
        "context_isolation_violation",
      );
    }
    if (candidate.status === "superseded" || candidate.status === "rejected") {
      excludedSuperseded += 1;
      continue;
    }
    candidates.push(candidate);
  }

  if (options.skillLifecycle) {
    for (const metadata of options.skillLifecycle.list("promoted")) {
      try {
        options.skillLifecycle.verify("promoted", metadata.id);
        candidates.push(parseContextItem({
          id: `skill_${metadata.id}`,
          kind: "skill",
          project: request.project,
          summary: metadata.description ? `${metadata.name}: ${metadata.description}` : metadata.name,
          provenance: {
            sourceType: "git",
            sourceRef: `skill:${metadata.id}@${metadata.contentSha256}`,
          },
          authority: "instruction",
          confidence: "verified",
          status: "current",
        }));
      } catch {
        // Tampered or invalid promoted skill: omit by hash check
      }
    }
  }

  const kindRank = new Map<string, number>();
  const requestedKinds = request.includeKinds ?? policy.kinds;
  requestedKinds.forEach((kind, index) => kindRank.set(kind, index));
  const kindPreference = (item: ContextItem): number => {
    const rank = kindRank.get(item.kind);
    return rank === undefined ? requestedKinds.length : rank;
  };

  // Eligibility: open contradictions always compete; everything else must be
  // a requested kind and not an inert proposal (a non-current state or a
  // proposed skill), so proposals never consume budget.
  const allowedKinds = new Set<ContextItemKind>(requestedKinds);
  const inertProposal = (item: ContextItem): boolean =>
    (item.kind === "state" && item.status !== undefined && item.status !== "current")
    || (item.kind === "skill" && item.status === "proposed");
  const eligible = candidates.filter((item) =>
    contradictions.has(item.id) || (allowedKinds.has(item.kind) && !inertProposal(item)));

  const scoreList = scoreRelevance(request.task, eligible.map(contextItemRelevanceText));
  const scores = new Map<ContextItem, number>();
  eligible.forEach((item, index) => scores.set(item, scoreList[index] ?? 0));
  const scoreOf = (item: ContextItem): number => scores.get(item) ?? 0;
  // Recency from the item's own timestamps; never reads the clock.
  const when = (item: ContextItem): number => {
    const parsed = Date.parse(item.observedAt ?? item.validFrom ?? "");
    return Number.isFinite(parsed) ? parsed : 0;
  };

  const ordered = [...eligible].sort((left, right) =>
    (contradictions.has(right.id) ? 1 : 0) - (contradictions.has(left.id) ? 1 : 0)
    || (left.project === request.project ? 0 : 1) - (right.project === request.project ? 0 : 1)
    || (scoreOf(right) > 0 ? 1 : 0) - (scoreOf(left) > 0 ? 1 : 0)
    || kindPreference(left) - kindPreference(right)
    || scoreOf(right) - scoreOf(left)
    || CONFIDENCE_RANK[right.confidence] - CONFIDENCE_RANK[left.confidence]
    || AUTHORITY_WEIGHT[right.authority] - AUTHORITY_WEIGHT[left.authority]
    || when(right) - when(left)
    || compareCodeUnitIds(left.id, right.id),
  );

  // First-fit: an item that does not fit is deferred and smaller items keep
  // filling the budget.
  const selected: ContextItem[] = [];
  const unresolvedGaps: string[] = [];
  let characters = 0;
  let deferredForBudget = 0;
  for (const item of ordered) {
    if (selected.length >= MAX_CONTEXT_ITEMS) {
      unresolvedGaps.push("context item limit reached; refine the task or kinds");
      break;
    }
    const itemCharacters = contextItemCharacters(item);
    if (Math.ceil((characters + itemCharacters) / 4) > budget) {
      deferredForBudget += 1;
      continue;
    }
    characters += itemCharacters;
    selected.push(item);
  }
  if (deferredForBudget > 0) {
    unresolvedGaps.push(selected.length === 0
      ? `budget of ${budget} tokens cannot fit any selected context`
      : `budget of ${budget} tokens reached; ${deferredForBudget} candidates deferred`);
  }
  if (candidates.length === 0) {
    unresolvedGaps.push("no context records exist for this project yet");
  }

  const bySection = (kind: ContextItemKind): ContextItem[] =>
    selected.filter((item) => item.kind === kind && !contradictions.has(item.id));

  const packet: ContextPacket = {
    workingState: options.workingState ?? {},
    currentState: bySection("state").filter((item) => item.status === "current" || item.status === undefined),
    knowledge: bySection("knowledge"),
    evidence: bySection("evidence"),
    episodes: bySection("episode"),
    skills: bySection("skill").filter((item) => item.status !== "proposed"),
    contradictions: selected.filter((item) => contradictions.has(item.id)),
    unresolvedGaps,
    provenanceSummary: provenanceSummaryOf(selected),
    estimatedTokens: estimateContextTokens(selected),
  };
  validateContextPacketContents(request, packet);

  return {
    packet,
    audit: {
      request: {
        project: request.project,
        role: request.role,
        task: request.task,
        ...(request.workflowRunId !== undefined ? { workflowRunId: request.workflowRunId } : {}),
        ...(request.stageId !== undefined ? { stageId: request.stageId } : {}),
      },
      selectedIds: selected.map((item) => item.id),
      provenanceSummary: packet.provenanceSummary,
      estimatedTokens: packet.estimatedTokens,
      budgetTokens: budget,
      candidateCount: candidates.length,
      excludedSuperseded,
      unresolvedGaps,
      relevance: {
        taskTokens: new Set(relevanceTokens(request.task)).size,
        matchedCandidates: scoreList.filter((score) => score > 0).length,
        selected: selected.map((item) => roundRelevance(scoreOf(item))),
      },
    },
  };
}

/** Convert a durable journal entry into a pool context item. Journal content
 * is evidence about work, never policy: entries become `evidence` or
 * `knowledge` items with workflow provenance and the entry as sourceRef. */
export function journalEntryToContextItem(entry: WorkflowJournalEntry, project: string): ContextItem {
  const kind: ContextItemKind = entry.category === "plan" || entry.category === "decision" || entry.category === "lesson"
    ? "knowledge"
    : entry.category === "skill-candidate"
      ? "skill"
      : "evidence";
  const item: ContextItem = {
    id: `journal_${entry.id}`,
    kind,
    project,
    summary: entry.summary,
    provenance: {
      sourceType: "workflow",
      sourceRef: `journal:${entry.id}`,
    },
    authority: entry.category === "decision" || entry.category === "plan" ? "evidence" : "evidence",
    confidence: entry.severity === "error" ? "probable" : "probable",
    observedAt: entry.createdAt,
    evidenceRefs: entry.evidence
      .filter((ref) => ref.length > 0 && ref.length <= 200)
      .slice(0, 16),
  };
  if (kind === "skill") item.status = "proposed";
  return parseContextItem(item);
}

/** Convert an authored memory record into a pool context item. Git-authored
 * records have instruction authority. Records scoped to 'operator' are placed
 * in the '_shared' project namespace so they can be consumed cross-project
 * as shared defaults without violating project boundaries. */
export function memoryRecordToContextItem(record: MemoryRecord, project: string): ContextItem {
  let authority: ContextAuthority = "instruction";
  if (record.authority === "evidence") {
    authority = "evidence";
  }

  const sourceType: ContextSourceType = CONTEXT_SOURCE_TYPES.includes(record.provenance?.sourceType as ContextSourceType)
    ? (record.provenance.sourceType as ContextSourceType)
    : "git";

  const targetProject = record.scope === "operator" ? "_shared" : project;

  return parseContextItem({
    id: `mem_${record.id}`,
    kind: "knowledge",
    scope: record.scope,
    project: targetProject,
    summary: record.summary,
    provenance: {
      sourceType,
      sourceRef: record.provenance?.sourceRef ?? `memory:${record.id}.md`,
    },
    authority,
    confidence: record.confidence,
    status: record.lifecycle === "active" ? "current" : "superseded",
    evidenceRefs: record.evidenceRefs && record.evidenceRefs.length > 0 ? record.evidenceRefs : undefined,
  });
}

/** Which pool records support a claim: the item itself plus its full
 * derivation lineage and evidence refs, metadata only (no raw bodies). */
export function explainContextItem(id: string, pool: ContextItem[]): {
  item: ContextItem | undefined;
  lineage: string[];
  evidenceRefs: string[];
  sources: { id: string; sourceType: string; sourceRef?: string }[];
} {
  const byId = new Map(pool.map((item) => [item.id, item]));
  const item = byId.get(id);
  if (!item) return { item: undefined, lineage: [], evidenceRefs: [], sources: [] };
  const lineage: string[] = [];
  const queue = [...(item.provenance.derivedFrom ?? [])];
  const seen = new Set<string>();
  while (queue.length > 0) {
    const ancestorId = queue.shift()!;
    if (seen.has(ancestorId)) continue;
    seen.add(ancestorId);
    lineage.push(ancestorId);
    const ancestor = byId.get(ancestorId);
    for (const older of ancestor?.provenance.derivedFrom ?? []) queue.push(older);
  }
  lineage.sort();
  const sources = [item, ...lineage.map((ancestorId) => byId.get(ancestorId))]
    .filter((candidate): candidate is ContextItem => candidate !== undefined)
    .map((candidate) => ({
      id: candidate.id,
      sourceType: candidate.provenance.sourceType,
      ...(candidate.provenance.sourceRef !== undefined ? { sourceRef: candidate.provenance.sourceRef } : {}),
    }));
  return {
    item,
    lineage,
    evidenceRefs: item.evidenceRefs ?? [],
    sources,
  };
}