feat(323): P0 gates measured + P1 — a model on this computer runs an agent

apps/aflow-executor-ai/src/handlers/agentTurnTimeoutSpec.test.ts

@@ -4,9 +4,19 @@ import { AiHandler } from './ai/AiHandler.js';
 
 const handler = new AiHandler({ payloadStore: {} as never });
 
-function ctx(operationId: string, executionTimeoutMs?: number): ExecutorContext {
+function ctx(
+  operationId: string,
+  executionTimeoutMs?: number,
+  input: Record<string, unknown> = {},
+): ExecutorContext {
   return {
     operationId,
+    job: {
+      stepExecutionId: `step-${operationId}-${String(input['model'])}`,
+      attempt: 1,
+      inputRef: 'inline:e30=',
+    },
+    readPayload: () => Promise.resolve(input),
     stepDefinition:
       executionTimeoutMs !== undefined ? { timeout: { executionTimeoutMs } } : undefined,
   } as unknown as ExecutorContext;
@@ -32,4 +42,31 @@ describe('AiHandler.resolveTimeoutMs', () => {
   it('non-agent operations keep the executor default', async () => {
     expect(await handler.resolveTimeoutMs(ctx('ai.text.generate'))).toBeUndefined();
   });
+
+  it('a model on this computer may read its prompt for minutes before its first token', async () => {
+    const cloud = (await handler.resolveTimeoutMs(
+      ctx('ai.agent.turn', undefined, { model: 'glm-pro' }),
+    )) as {
+      idleMs: number;
+    };
+    const local = (await handler.resolveTimeoutMs(
+      ctx('ai.text.generate', undefined, { model: 'gemma-local' }),
+    )) as { idleMs: number; maxMs: number };
+    expect(local.idleMs).toBeGreaterThan(cloud.idleMs);
+    expect(local.maxMs).toBeGreaterThan(local.idleMs);
+  });
+});
+
+describe('AiHandler.admit', () => {
+  it('asks nothing of a step whose model is served elsewhere', async () => {
+    const withRedis = new AiHandler({ payloadStore: {} as never, redis: {} as never });
+    const waiting = {
+      signal: new AbortController().signal,
+      waiting: () => undefined,
+      stillWaiting: () => undefined,
+    };
+    expect(
+      await withRedis.admit(ctx('ai.agent.turn', undefined, { model: 'glm-pro' }), waiting),
+    ).toBeUndefined();
+  });
 });

apps/aflow-executor-ai/src/handlers/ai/AiHandler.ts

@@ -4,6 +4,8 @@ import type {
   ExecutorContext,
   TimeoutSpec,
   StepResult,
+  StepAdmission,
+  StepAdmissionWaiting,
 } from '@aflow/executor-runtime';
 import { failureWithError, validationError, internalError } from '@aflow/executor-runtime';
 import {
@@ -29,10 +31,11 @@ import { AIClientError } from '@aflow/ai-client';
 import { positiveMsEnv } from '@aflow/lib';
 import type { ErrorObject } from 'ajv';
 import type { PayloadStore } from '@aflow/payload-store';
-import type { Redis } from '@aflow/redis';
+import { acquireRuntimeLease, type Redis } from '@aflow/redis';
 import type { PostgresJsDatabase } from 'drizzle-orm/postgres-js';
 import { createAsyncJobRepository, createTenantContext } from '@aflow/database';
-import { AISetupError, checkBudgetMock } from './aiClient.js';
+import { AISetupError, checkBudgetMock, resolveLocalRuntime } from './aiClient.js';
+import { inferProviderForModelRef } from '@aflow/ai-client';
 import { ajv } from './ajv.js';
 import {
   AiGenerateInputSchema,
@@ -83,28 +86,104 @@ type OperationHandler = (ctx: ExecutorContext, input: unknown) => Promise<StepRe
 // the 20-minute ceiling is what bounds a genuinely stuck turn.
 const AGENT_TURN_STREAM_IDLE_MS = positiveMsEnv('AI_STREAM_IDLE_TIMEOUT_MS', 300_000);
 const AGENT_TURN_MAX_MS = positiveMsEnv('AI_AGENT_TURN_MAX_TIMEOUT_MS', 1_200_000);
+/**
+ * A model on the operator's machine reads its whole prompt before it writes
+ * anything, and nothing streams meanwhile: at 250 tokens a second a 75K-token
+ * prompt outlasts the cloud idle window (Plan 323 §9 G5). Its window is the
+ * prefill a long prompt takes on a laptop.
+ */
+const LOCAL_MODEL_IDLE_MS = positiveMsEnv('AI_LOCAL_MODEL_IDLE_TIMEOUT_MS', 900_000);
+
+/** Attempts whose input admission read and whose handler has not yet run; bounded by the executor's concurrency. */
+const MAX_CACHED_INPUTS = 64;
+
+function stepModel(input: unknown): string | undefined {
+  const model = (input as { model?: unknown } | null)?.model;
+  return typeof model === 'string' && model.length > 0 ? model : undefined;
+}
+
+function attemptKey(ctx: ExecutorContext): string {
+  return `${ctx.job.stepExecutionId}:${ctx.job.attempt}`;
+}
 
 export class AiHandler implements StepHandler {
   readonly stepType = 'ai';
 
-  resolveTimeoutMs(ctx: ExecutorContext): Promise<TimeoutSpec | undefined> {
+  async resolveTimeoutMs(ctx: ExecutorContext): Promise<TimeoutSpec | undefined> {
     // An explicit step-definition timeout wins in both directions and stays a
     // flat wall clock (an operator who set a number meant a number); this only
     // replaces the executor-wide flat default.
     if (ctx.stepDefinition?.timeout?.executionTimeoutMs !== undefined) {
-      return Promise.resolve(undefined);
+      return undefined;
+    }
+    const model = stepModel(await this.stepInput(ctx));
+    if (model !== undefined && inferProviderForModelRef(model) === 'ollama') {
+      return { idleMs: LOCAL_MODEL_IDLE_MS, maxMs: AGENT_TURN_MAX_MS };
     }
     if (ctx.operationId === 'ai.agent.turn') {
-      return Promise.resolve({ idleMs: AGENT_TURN_STREAM_IDLE_MS, maxMs: AGENT_TURN_MAX_MS });
+      return { idleMs: AGENT_TURN_STREAM_IDLE_MS, maxMs: AGENT_TURN_MAX_MS };
     }
     // A render outlasts the executor-wide default, and the step that dies
     // waiting was still billed for it.
     if (ASYNC_JOB_LIFECYCLE_OPERATIONS.has(ctx.operationId)) {
-      return Promise.resolve(resolveVideoBudget().stepBudgetMs);
+      return resolveVideoBudget().stepBudgetMs;
+    }
+    return undefined;
+  }
+
+  /**
+   * A step whose model runs on the operator's machine waits for that runtime
+   * to have room (Plan 323 D7) — before its timeout starts, and without its
+   * executor slot.
+   */
+  async admit(
+    ctx: ExecutorContext,
+    waiting: StepAdmissionWaiting,
+  ): Promise<StepAdmission | undefined> {
+    if (!this.redis) return undefined;
+    const model = stepModel(await this.stepInput(ctx));
+    if (model === undefined) return undefined;
+    const runtime = await resolveLocalRuntime(ctx, model);
+    if (!runtime) return undefined;
+    const lease = await acquireRuntimeLease({
+      redis: this.redis,
+      runtimeId: runtime.runtimeId,
+      limit: runtime.maxParallel,
+      holderId: attemptKey(ctx),
+      wait: true,
+      signal: waiting.signal,
+      onWaiting: () => {
+        waiting.waiting({ runtimeId: runtime.runtimeId, limit: runtime.maxParallel, model });
+      },
+      onStillWaiting: waiting.stillWaiting,
+    });
+    return lease ?? undefined;
+  }
+
+  /**
+   * The step's input, read once per attempt: admission and the timeout read
+   * its model before the handler runs, and an agent turn's input is large.
+   */
+  private stepInput(ctx: ExecutorContext): Promise<unknown> {
+    const key = attemptKey(ctx);
+    let input = this.inputs.get(key);
+    if (!input) {
+      input = ctx.readPayload(ctx.job.inputRef);
+      this.inputs.set(key, input);
+      // An attempt that never reaches `execute` (cancelled while it waited)
+      // must not hold its input for the life of the process.
+      if (this.inputs.size > MAX_CACHED_INPUTS) {
+        const oldest = this.inputs.keys().next().value;
+        if (oldest !== undefined) this.inputs.delete(oldest);
+      }
     }
-    return Promise.resolve(undefined);
+    return input;
   }
 
+  private readonly inputs = new Map<string, Promise<unknown>>();
+
+  private readonly redis: Redis | undefined;
+
   private readonly payloadStore: PayloadStore;
 
   private readonly operations: Record<
@@ -117,6 +196,7 @@ export class AiHandler implements StepHandler {
 
   constructor(deps: { payloadStore: PayloadStore; db?: PostgresJsDatabase; redis?: Redis }) {
     this.payloadStore = deps.payloadStore;
+    this.redis = deps.redis;
 
     const db = deps.db;
     const depsForHandlers: HandlerDeps = {
@@ -327,7 +407,8 @@ export class AiHandler implements StepHandler {
 
   async execute(ctx: ExecutorContext): Promise<StepResult> {
     const operationId = ctx.operationId;
-    const input = await ctx.readPayload(ctx.job.inputRef);
+    const input = await this.stepInput(ctx);
+    this.inputs.delete(attemptKey(ctx));
     const typedInput = input as Record<string, unknown>;
 
     const budgetResult = checkBudgetMock({

apps/aflow-executor-ai/src/handlers/ai/aiClient.ts

@@ -2,15 +2,17 @@ import {
   createAIClient,
   inferProviderForModelRef,
   type AIClient,
-  type ProviderConfig,
   type BudgetCheckResult,
   type AIProvider,
 } from '@aflow/ai-client';
 import type { ExecutorContext } from '@aflow/executor-runtime';
-import type { AflowError } from '@aflow/schemas';
+import { localRuntimeId, type AflowError } from '@aflow/schemas';
 import {
   CredentialResolver,
   credentialMissingMessage,
+  localRuntimeSettings,
+  missingRequiredSecrets,
+  providerConfigFromCredential,
   type CredentialContext,
   type CredentialLoader,
 } from '@aflow/credential-resolver';
@@ -27,6 +29,7 @@ const AI_TO_CREDENTIAL_PROVIDER: Record<string, string> = {
   xai: 'xai',
   runware: 'runware',
   typesafe: 'typesafe',
+  ollama: 'ollama',
 };
 
 // ---------------------------------------------------------------------------
@@ -284,10 +287,10 @@ export async function getAIClientForContext(
     );
   }
 
-  const apiKey = resolved.secrets['api_key'];
-  if (!apiKey) {
+  const missingSecrets = missingRequiredSecrets(credentialProviderId, resolved);
+  if (missingSecrets.length > 0) {
     throw new AISetupError(
-      `Credential for ${credentialProviderId} exists but has no api_key. ` +
+      `Credential for ${credentialProviderId} exists but has no ${missingSecrets.join(', ')}. ` +
         'Update your credentials in Settings → Credentials.',
       'AI_PROVIDER_API_KEY_MISSING',
       'configuration',
@@ -305,16 +308,10 @@ export async function getAIClientForContext(
   }
 
   // 5. Build single-provider AIClient
-  const providerConfig: ProviderConfig = { apiKey };
-  if (aiProvider === 'openai' && resolved.config['org_id']) {
-    providerConfig.organization = resolved.config['org_id'];
-  }
-
-  const providers: Partial<Record<AIProvider, ProviderConfig>> = {
-    [aiProvider]: providerConfig,
-  };
-
-  const client = createAIClient({ providers, defaultProvider: aiProvider });
+  const client = createAIClient({
+    providers: { [aiProvider]: providerConfigFromCredential(credentialProviderId, resolved) },
+    defaultProvider: aiProvider,
+  });
 
   CLIENT_CACHE.set(cacheKey, {
     client,
@@ -323,3 +320,26 @@ export async function getAIClientForContext(
 
   return client;
 }
+
+/**
+ * The runtime on the operator's machine a model is served by, and how many
+ * requests it takes at once, or null for a model served anywhere else.
+ * A local model with no connection resolves to null too: the request then
+ * fails on the missing credential, which says what to do.
+ */
+export async function resolveLocalRuntime(
+  ctx: ExecutorContext,
+  model: string,
+): Promise<{ runtimeId: string; maxParallel: number } | null> {
+  if (inferProviderForModelRef(model) !== 'ollama') return null;
+  const { tenantId, credentialOwnerId, spaceId } = ctx.job;
+  if (!credentialResolver || !credentialOwnerId || !spaceId) return null;
+  const resolved = await credentialResolver
+    .resolve('ollama', { tenantId, credentialOwnerId, spaceId })
+    .catch(() => null);
+  if (!resolved) return null;
+  return {
+    runtimeId: localRuntimeId(providerConfigFromCredential('ollama', resolved).baseUrl),
+    maxParallel: localRuntimeSettings('ollama', resolved).maxParallel,
+  };
+}

apps/aflow-executor-ai/src/handlers/ai/handlers/agentTurn.test.ts

@@ -605,6 +605,8 @@ describe('handleAgentTurn (integration seams)', () => {
     };
 
     vi.mocked(getAIClientForContext).mockResolvedValue({
+      contextWindowFor: () => undefined,
+
       getModel: () => ({
         id: model,
         provider,
@@ -638,6 +640,8 @@ describe('handleAgentTurn (integration seams)', () => {
     };
 
     vi.mocked(getAIClientForContext).mockResolvedValue({
+      contextWindowFor: () => undefined,
+
       getModel: () => ({
         id: model,
         provider,
@@ -672,6 +676,8 @@ describe('handleAgentTurn (integration seams)', () => {
 
     let call = 0;
     vi.mocked(getAIClientForContext).mockResolvedValue({
+      contextWindowFor: () => undefined,
+
       getModel: () => ({
         id: model,
         provider,
@@ -719,6 +725,8 @@ describe('handleAgentTurn (integration seams)', () => {
 
     let call = 0;
     vi.mocked(getAIClientForContext).mockResolvedValue({
+      contextWindowFor: () => undefined,
+
       getModel: () => ({
         id: model,
         provider,
@@ -787,6 +795,8 @@ describe('handleAgentTurn (integration seams)', () => {
       usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 },
     });
     vi.mocked(getAIClientForContext).mockResolvedValue({
+      contextWindowFor: () => undefined,
+
       getModel: () => ({
         id: model,
         provider,
@@ -834,6 +844,8 @@ describe('handleAgentTurn (integration seams)', () => {
     }
 
     vi.mocked(getAIClientForContext).mockResolvedValue({
+      contextWindowFor: () => undefined,
+
       getModel: () => ({
         id: 'claude-3-test',
         provider: 'anthropic',
@@ -889,6 +901,8 @@ describe('handleAgentTurn (integration seams)', () => {
     const generateJson = vi.fn();
 
     vi.mocked(getAIClientForContext).mockResolvedValue({
+      contextWindowFor: () => undefined,
+
       getModel: () => ({
         id: 'gpt-6.1-sol',
         provider: 'openai',
@@ -938,6 +952,8 @@ describe('handleAgentTurn (integration seams)', () => {
     const generateJson = vi.fn();
 
     vi.mocked(getAIClientForContext).mockResolvedValue({
+      contextWindowFor: () => undefined,
+
       getModel: () => ({
         id: 'grok-4.7',
         provider: 'xai',
@@ -1002,6 +1018,8 @@ describe('handleAgentTurn (integration seams)', () => {
     const generateJson = vi.fn();
 
     vi.mocked(getAIClientForContext).mockResolvedValue({
+      contextWindowFor: () => undefined,
+
       getModel: () => ({
         id: 'gpt-6.1-sol',
         provider: 'openai',