refactor(context): decouple provider execution from context collector

src/agent/runtime/context-manager.ts

@@ -8,6 +8,7 @@ import {
   createDefaultTokenEstimator,
   PromptAssembler,
   RuleBasedContextCompressor,
+  type TokenEstimator,
   TranscriptProjector,
   type BuildContextDiagnostics,
   type ContextDocument,
@@ -36,27 +37,35 @@ export type ContextManagerOptions = {
   observer: AgentObserver;
   /** Agent-level context providers invoked for every run. */
   contextProviders?: readonly AgentContextProvider[];
+
+  /**
+   * Context build core dependencies.
+   */
+  transcriptProjector?: TranscriptProjector;
+
+  tokenEstimator?: TokenEstimator;
+
+  budgeter?: ContextBudgeter;
+
+  compressionPipeline?: ContextCompressionPipeline;
+
+  truncator?: ContextTruncator;
+
+  promptAssembler?: PromptAssembler;
+
+  systemPromptService?: SystemPromptService;
 };
 
 /** Builds model context and owns transcript mutation for assistant/tool turns. */
 export class ContextManager {
   private readonly collector: ContextCollector;
-  private readonly transcriptProjector = new TranscriptProjector();
-  private readonly tokenEstimator = createDefaultTokenEstimator();
-  private readonly budgeter = new ContextBudgeter({
-    tokenEstimator: this.tokenEstimator,
-  });
-  private readonly compressionPipeline = new ContextCompressionPipeline({
-    compressor: new RuleBasedContextCompressor({
-      tokenEstimator: this.tokenEstimator,
-    }),
-    tokenEstimator: this.tokenEstimator,
-  });
-  private readonly truncator = new ContextTruncator({
-    tokenEstimator: this.tokenEstimator,
-  });
-  private readonly promptAssembler = new PromptAssembler();
-  private readonly systemPromptService = new SystemPromptService();
+  private readonly transcriptProjector: TranscriptProjector;
+  private readonly tokenEstimator: TokenEstimator;
+  private readonly budgeter: ContextBudgeter;
+  private readonly compressionPipeline: ContextCompressionPipeline;
+  private readonly truncator: ContextTruncator;
+  private readonly promptAssembler: PromptAssembler;
+  private readonly systemPromptService: SystemPromptService;
   private readonly runStates = new WeakMap<AgentRunState, ContextManagerRunState>();
 
   /** Creates a context manager with static agent-level context providers. */
@@ -65,6 +74,36 @@ export class ContextManager {
       memory: options.memory,
       contextProviders: options.contextProviders,
     });
+
+    this.tokenEstimator = options.tokenEstimator ?? createDefaultTokenEstimator();
+
+    this.transcriptProjector = options.transcriptProjector ?? new TranscriptProjector();
+
+    this.budgeter =
+      options.budgeter ??
+      new ContextBudgeter({
+        tokenEstimator: this.tokenEstimator,
+      });
+
+    this.compressionPipeline =
+      options.compressionPipeline ??
+      new ContextCompressionPipeline({
+        compressor: new RuleBasedContextCompressor({
+          tokenEstimator: this.tokenEstimator,
+        }),
+
+        tokenEstimator: this.tokenEstimator,
+      });
+
+    this.truncator =
+      options.truncator ??
+      new ContextTruncator({
+        tokenEstimator: this.tokenEstimator,
+      });
+
+    this.promptAssembler = options.promptAssembler ?? new PromptAssembler();
+
+    this.systemPromptService = options.systemPromptService ?? new SystemPromptService();
   }
 
   /** Initializes static context sources and the mutable transcript for a run. */

src/agent/types.ts

@@ -75,8 +75,15 @@ export type AgentContextValue =
 /** Deferred context provider used to attach dynamic runtime context. */
 export type AgentContextProvider = {
   name: string;
-  description?: string;
-  build(context: AgentRunContext): AgentContextValue | Promise<AgentContextValue>;
+  critical?: boolean;
+  timeoutMs?: number;
+
+  build(
+    context: AgentRunContext,
+    options: {
+      signal: AbortSignal;
+    },
+  ): Promise<AgentContextValue>;
 };
 
 /** User-facing input for one agent run. */

src/context/collector/context-collector.ts

@@ -2,56 +2,109 @@ import type {AgentContextProvider, AgentContextValue} from "../../agent/types.js
 import type {AgentMemory} from "../../agent/runtime/memory.js";
 import type {AgentRunState} from "../../agent/runtime/run-state.js";
 import type {ContextDocument, ContextSection} from "../types.js";
+import {ContextProviderExecutor} from "./context-provider-executor.js";
 
 export type ContextCollectorOptions = {
   memory: AgentMemory;
+
+  /**
+   * Static context providers.
+   *
+   * These providers are registered when creating Agent runtime.
+   */
   contextProviders?: readonly AgentContextProvider[];
 };
 
 export type CollectContextOptions = {
   stage?: ContextDocument["metadata"]["stage"];
 };
 
-/** Collects run context sources without applying ordering or budget policy. */
+/**
+ * Collects context sources and builds ContextDocument.
+ *
+ * Responsibilities:
+ * - load memory
+ * - collect workspace/user context
+ * - execute context providers
+ * - convert values into ContextSection
+ *
+ * Not responsible for:
+ * - token budgeting
+ * - compression
+ * - ordering policy
+ * - prompt assembling
+ */
 export class ContextCollector {
-  private readonly contextProviders: AgentContextProvider[];
+  private readonly contextProviders: readonly AgentContextProvider[];
+
+  private readonly providerExecutor: ContextProviderExecutor;
 
   constructor(private readonly options: ContextCollectorOptions) {
-    this.contextProviders = [...(options.contextProviders ?? [])];
+    this.contextProviders = options.contextProviders ?? [];
+
+    this.providerExecutor = new ContextProviderExecutor();
   }
 
   async collect(
     run: AgentRunState,
     options: CollectContextOptions = {},
   ): Promise<ContextDocument> {
-    const sections = [
-      ...(await this.loadMemory(run)).map((value) => toContextSection(value, "memory")),
+    const sections: ContextSection[] = [];
+
+    /**
+     * Memory Context
+     */
+    const memories = await this.loadMemory(run);
+
+    sections.push(...memories.map((value) => toContextSection(value, "memory")));
+
+    /**
+     * User provided context
+     */
+    sections.push(
       ...(run.input.context ?? []).map((value) => toContextSection(value, "user")),
-      {
-        title: "Workspace Profile",
-        priority: 100,
-        content: JSON.stringify(run.workspaceProfile, null, 2),
-        source: {kind: "workspace"},
-      } satisfies ContextSection,
-    ];
+    );
+
+    /**
+     * Workspace Context
+     */
+    sections.push({
+      title: "Workspace Profile",
+
+      priority: 100,
 
+      content: JSON.stringify(run.workspaceProfile, null, 2),
+
+      source: {
+        kind: "workspace",
+      },
+    });
+
+    /**
+     * Context Providers
+     */
     const providers = [...this.contextProviders, ...(run.input.contextProviders ?? [])];
 
-    for (const provider of providers) {
-      const value = await provider.build(run.context);
-      sections.push(
-        typeof value === "string"
-          ? {
-              title: provider.name,
-              content: value,
-              source: {kind: "provider", ref: provider.name},
-            }
-          : {
-              title: value.title ?? provider.name,
-              ...value,
-              source: {kind: "provider", ref: provider.name},
-            },
-      );
+    const providerResults = await this.providerExecutor.execute(providers, run.context);
+
+    for (const result of providerResults) {
+      /**
+       * Non critical provider failure
+       *
+       * Executor already isolates errors.
+       * Collector only skips invalid results.
+       */
+      if (!result.success) {
+        console.warn(`[ContextProvider] ${result.provider.name} failed:`, result.error);
+
+        continue;
+      }
+
+      if (result.value === undefined) {
+        continue;
+      }
+
+      sections.push(toProviderSection(result.provider, result.value));
     }
 
     return {
@@ -68,17 +121,67 @@ export class ContextCollector {
   private async loadMemory(run: AgentRunState): Promise<AgentContextValue[]> {
     const runMemory = await this.options.memory.loadRunMemory?.(run);
     const projectMemory = await this.options.memory.loadProjectMemory?.(run);
+
     return [...(projectMemory ?? []), ...(runMemory ?? [])];
   }
 }
 
+/**
+ * Convert common context value to ContextSection.
+ */
 function toContextSection(
   value: AgentContextValue,
   source: "memory" | "user",
 ): ContextSection {
   if (typeof value === "string") {
-    return {content: value, source: {kind: source}};
+    return {
+      content: value,
+
+      source: {
+        kind: source,
+      },
+    };
+  }
+
+  return {
+    ...value,
+
+    source: {
+      kind: source,
+    },
+  };
+}
+
+/**
+ * Convert provider result to ContextSection.
+ */
+function toProviderSection(
+  provider: AgentContextProvider,
+  value: AgentContextValue,
+): ContextSection {
+  if (typeof value === "string") {
+    return {
+      title: provider.name,
+
+      content: value,
+
+      source: {
+        kind: "provider",
+
+        ref: provider.name,
+      },
+    };
   }
 
-  return {...value, source: {kind: source}};
+  return {
+    title: value.title ?? provider.name,
+
+    ...value,
+
+    source: {
+      kind: "provider",
+
+      ref: provider.name,
+    },
+  };
 }