@stackbilt/aegis-core 0.6.5 → 0.8.0

package/src/composite.ts

@@ -7,9 +7,10 @@
 //   Claude Sonnet: $3/$15 per MTok (synthesize)
 //   Expected total: $0.01-0.03 per composite query
 
+import type { LLMMessage, ToolResult as LLMToolResult } from '@stackbilt/llm-providers';
 import { askGroqJson } from './groq.js';
 import { buildContext, handleInProcessTool, callMcpWithRetry, resolveMcpTool } from './claude.js';
-import { toOpenAiTools, extractText, extractToolCalls, extractUsage, type AiChatResponse } from './workers-ai-chat.js';
+import { toOpenAiTools } from './workers-ai-chat.js';
 import { McpClient, McpRegistry } from './mcp-client.js';
 import { operatorConfig } from './operator/index.js';
 import { buildPersonaPreamble } from './operator/prompt-builder.js';
@@ -17,6 +18,7 @@ import { getCognitiveState, formatCognitiveContext } from './kernel/cognition.js
 import { getAttachedBlocks, assembleBlockContext } from './kernel/memory/blocks.js';
 import { getConversationHistory, budgetConversationHistory } from './kernel/memory/index.js';
 import { classifyCourtCard, type CourtCard, type CourtCardProfile } from './kernel/court-cards.js';
+import { buildLLMProviderFactory } from './kernel/provider-factory.js';
 import type { KernelIntent } from './kernel/types.js';
 import type { EdgeEnv } from './kernel/dispatch.js';
 
@@ -43,10 +45,6 @@ interface SubtaskResult {
 // ─── Cost rates ─────────────────────────────────────────────
 
 const GROQ_GPT_OSS_RATES = { input: 0.15, output: 0.60 };
-const CF_GPT_OSS_RATES = { input: 0.35, output: 0.75 };
-const CLAUDE_SONNET_RATES = { input: 3, output: 15 };
-const CLAUDE_OPUS_RATES = { input: 15, output: 75 };
-
 // ─── Phase 1: Orchestrate ───────────────────────────────────
 
 const ORCHESTRATOR_SYSTEM = `You are a task decomposition engine. Given a user query and a list of available tools, decompose the query into subtasks that can be executed independently.
@@ -160,13 +158,7 @@ function validateDagIntent(userQuery: string, dag: ExecutionDAG): boolean {
   return true;
 }
 
-// ─── Phase 2: Gather (CF Workers AI tool loop) ─────────────
-
-type ChatMessage =
-  | { role: 'system'; content: string }
-  | { role: 'user'; content: string }
-  | { role: 'assistant'; content: string; tool_calls?: Array<{ id: string; type: string; function: { name: string; arguments: string } }> }
-  | { role: 'tool'; tool_call_id: string; content: string };
+// ─── Phase 2: Gather (provider-backed Workers AI tool loop) ──
 
 const MAX_GATHER_ROUNDS = 6;
 
@@ -210,7 +202,6 @@ async function synthesize(
   courtCard?: CourtCardProfile,
 ): Promise<{ text: string; cost: number }> {
   const model = useOpus ? env.opusModel : env.claudeModel;
-  const rates = useOpus ? CLAUDE_OPUS_RATES : CLAUDE_SONNET_RATES;
 
   const subtaskSummary = subtaskResults.map(r => {
     // Include raw gathered data so synthesis can recover structured values the analysis step may have dropped
@@ -233,38 +224,17 @@ async function synthesize(
     }
   } catch { /* non-fatal — synthesize without cognitive context */ }
 
-  const response = await fetch(`${env.anthropicBaseUrl}/v1/messages`, {
-    method: 'POST',
-    headers: {
-      'Content-Type': 'application/json',
-      'x-api-key': env.anthropicApiKey,
-      'anthropic-version': '2023-06-01',
-    },
-    body: JSON.stringify({
-      model,
-      max_tokens: 4096,
-      system: `${buildPersonaPreamble()} Synthesize the analyzed subtask results into a coherent, actionable answer. Speak as AEGIS — the co-founder who knows the business inside-out. Be thorough but concise. Reference specific products, numbers, and context from the portfolio below. Never give generic consultant advice; give the answer a co-founder would give.${courtCard ? ` ${courtCard.synthesisVoice}` : ''}${contextSuffix}`,
-      messages: [{
-        role: 'user',
-        content: `Original query: ${intent.raw}\n\nSynthesis instruction: ${synthesisInstruction}\n\nSubtask results:\n${subtaskSummary}`,
-      }],
-    }),
+  const result = await buildLLMProviderFactory(env).generateResponse({
+    model,
+    maxTokens: 4096,
+    systemPrompt: `${buildPersonaPreamble()} Synthesize the analyzed subtask results into a coherent, actionable answer. Speak as AEGIS — the co-founder who knows the business inside-out. Be thorough but concise. Reference specific products, numbers, and context from the portfolio below. Never give generic consultant advice; give the answer a co-founder would give.${courtCard ? ` ${courtCard.synthesisVoice}` : ''}${contextSuffix}`,
+    messages: [{
+      role: 'user',
+      content: `Original query: ${intent.raw}\n\nSynthesis instruction: ${synthesisInstruction}\n\nSubtask results:\n${subtaskSummary}`,
+    }],
   });
 
-  if (!response.ok) {
-    const err = await response.text();
-    throw new Error(`Anthropic API error ${response.status}: ${err}`);
-  }
-
-  const data = await response.json<{
-    content: Array<{ type: string; text?: string }>;
-    usage: { input_tokens: number; output_tokens: number };
-  }>();
-
-  const text = data.content.filter(b => b.type === 'text').map(b => b.text ?? '').join('');
-  const cost = (data.usage.input_tokens * rates.input + data.usage.output_tokens * rates.output) / 1_000_000;
-
-  return { text: text || '(no synthesis)', cost };
+  return { text: result.message || '(no synthesis)', cost: result.usage.cost };
 }
 
 // ─── Groq synthesis fallback ────────────────────────────────
@@ -363,7 +333,10 @@ export async function executeComposite(
     githubRepo: env.githubRepo,
     braveApiKey: env.braveApiKey,
     roundtableDb: env.roundtableDb,
+    memoryBinding: env.memoryBinding,
+    resendApiKeys: { resendApiKey: env.resendApiKey, resendApiKeyPersonal: env.resendApiKeyPersonal },
     userQuery: intent.raw,
+    edgeEnv: env,
   }, env.roundtableDb);
 
   // Load conversation history for context continuity
@@ -602,7 +575,7 @@ async function gatherSubtaskInstrumented(
     ? `Original request: ${originalQuery}\n\nYour subtask: ${subtask.description}\n\nIMPORTANT: Use exact identifiers (UUIDs, IDs, enum values) from the original request when calling tools.`
     : subtask.description;
 
-  const messages: ChatMessage[] = [
+  const messages: LLMMessage[] = [
     { role: 'system', content: `${systemPrompt}\n\nFocus: ${subtask.description}\nGather the data needed and return your findings.` },
     { role: 'user', content: userContent },
   ];
@@ -613,29 +586,26 @@ async function gatherSubtaskInstrumented(
   // Tool loop — up to MAX_GATHER_ROUNDS
   for (let round = 0; round < MAX_GATHER_ROUNDS; round++) {
     subrequestCount += 1; // 1 AI call per round
-    const result = await env.ai.run(env.gptOssModel as Parameters<Ai['run']>[0], {
-      messages,
-      ...(openAiTools.length > 0 ? { tools: openAiTools } : {}),
-      max_tokens: 2048,
+    const result = await buildLLMProviderFactory(env).generateResponse({
+      messages: [...messages],
+      model: env.gptOssModel,
+      ...(openAiTools.length > 0 ? { tools: openAiTools as Parameters<ReturnType<typeof buildLLMProviderFactory>['generateResponse']>[0]['tools'] } : {}),
+      maxTokens: 2048,
       temperature: 0.2,
-      top_p: 0.9,
-      frequency_penalty: 0.3,
-    } as Record<string, unknown>) as AiChatResponse;
-
-    const usage = extractUsage(result);
-    if (usage) {
-      totalCost += (usage.prompt_tokens * CF_GPT_OSS_RATES.input
-        + usage.completion_tokens * CF_GPT_OSS_RATES.output) / 1_000_000;
-    }
+      topP: 0.9,
+      frequencyPenalty: 0.3,
+    });
 
-    const toolCalls = extractToolCalls(result);
-    const responseText = extractText(result);
+    totalCost += result.usage.cost;
+    const toolCalls = result.toolCalls ?? [];
+    const responseText = result.message;
 
     if (toolCalls.length === 0) {
       return { gathered: responseText ?? '(no data gathered)', cost: totalCost, subrequestCount };
     }
 
-    messages.push({ role: 'assistant', content: responseText ?? '', tool_calls: toolCalls });
+    messages.push({ role: 'assistant', content: responseText ?? '', toolCalls });
+    const toolResults: LLMToolResult[] = [];
 
     for (const call of toolCalls) {
       subrequestCount += 1; // 1 subrequest per tool call (external fetch or DB query)
@@ -650,6 +620,7 @@ async function gatherSubtaskInstrumented(
         { apiKey: env.anthropicApiKey, model: env.claudeModel, baseUrl: env.anthropicBaseUrl },
         env.memoryBinding,
         { resendApiKey: env.resendApiKey, resendApiKeyPersonal: env.resendApiKeyPersonal },
+        env,
       );
 
       if (inProcess !== null) {
@@ -663,26 +634,29 @@ async function gatherSubtaskInstrumented(
         }
       }
 
-      messages.push({ role: 'tool', tool_call_id: call.id, content: toolResult });
+      toolResults.push({ id: call.id, output: toolResult });
     }
+    messages.push({ role: 'user', content: '', toolResults });
   }
 
   // Force summary if tool rounds exhausted
   // Condense messages: strip tool_calls metadata and truncate tool results
   // to prevent context overflow when sending to GPT-OSS without tools definition
-  const condensedGather: ChatMessage[] = [messages[0]]; // system
+  const condensedGather: LLMMessage[] = [messages[0]]; // system
   const gatherFindings: string[] = [];
   for (let i = 1; i < messages.length; i++) {
     const msg = messages[i];
-    if (msg.role === 'user' && !('tool_call_id' in msg)) {
-      condensedGather.push(msg);
+    if (msg.role === 'user' && !msg.toolResults) {
+      condensedGather.push({ role: 'user', content: msg.content });
     } else if (msg.role === 'assistant' && msg.content) {
       gatherFindings.push(msg.content);
-    } else if (msg.role === 'tool') {
-      const truncated = msg.content.length > 2000
-        ? msg.content.slice(0, 2000) + '... [truncated]'
-        : msg.content;
-      gatherFindings.push(truncated);
+    } else if (msg.toolResults) {
+      for (const toolResult of msg.toolResults) {
+        const truncated = toolResult.output.length > 2000
+          ? toolResult.output.slice(0, 2000) + '... [truncated]'
+          : toolResult.output;
+        gatherFindings.push(truncated);
+      }
     }
   }
   if (gatherFindings.length > 0) {
@@ -699,17 +673,13 @@ async function gatherSubtaskInstrumented(
   condensedGather.push({ role: 'user', content: 'Summarize all data gathered so far. Return the raw findings.' });
 
   subrequestCount += 1;
-  const summaryResult = await env.ai.run(env.gptOssModel as Parameters<Ai['run']>[0], {
+  const summaryResult = await buildLLMProviderFactory(env).generateResponse({
     messages: condensedGather,
-    max_tokens: 2048,
+    model: env.gptOssModel,
+    maxTokens: 2048,
     temperature: 0.2,
-  } as Record<string, unknown>) as AiChatResponse;
-
-  const summaryUsage = extractUsage(summaryResult);
-  if (summaryUsage) {
-    totalCost += (summaryUsage.prompt_tokens * CF_GPT_OSS_RATES.input
-      + summaryUsage.completion_tokens * CF_GPT_OSS_RATES.output) / 1_000_000;
-  }
+  });
 
-  return { gathered: extractText(summaryResult) ?? '(gather exhausted)', cost: totalCost, subrequestCount };
+  totalCost += summaryResult.usage.cost;
+  return { gathered: summaryResult.message ?? '(gather exhausted)', cost: totalCost, subrequestCount };
 }

package/src/core.ts

@@ -27,6 +27,7 @@ import { observability } from './routes/observability.js';
 import { pages } from './routes/pages.js';
 import { ccTasks } from './routes/cc-tasks.js';
 import { messages } from './routes/messages.js';
+import { chatWs } from './routes/chat-ws.js';
 import { dynamicToolsRoutes } from './routes/dynamic-tools.js';
 
 // ─── Scheduled Task Plugin ──────────────────────────────────
@@ -213,6 +214,7 @@ export function createAegisApp(config: AegisAppConfig): AegisApp {
   app.route('/', pages);
   app.route('/', ccTasks);
   app.route('/', messages);
+  app.route('/', chatWs);
   app.route('/', dynamicToolsRoutes);
 
   // ── Extension routes ──
@@ -268,6 +270,7 @@ export function createAegisApp(config: AegisAppConfig): AegisApp {
     pages,
     ccTasks,
     messages,
+    chatWs,
     dynamicTools: dynamicToolsRoutes,
   };
 
@@ -288,6 +291,8 @@ export type {
   MessageMetadata,
 } from './types.js';
 
+export { ChatSession } from './durable-objects/chat-session.js';
+
 export type { EdgeEnv } from './kernel/dispatch.js';
 
 export type {

package/src/durable-objects/chat-session-auth.ts

@@ -0,0 +1,20 @@
+const UUID_RE = /^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i;
+
+export type ConversationOwnership = 'owned' | 'not_owned' | 'not_found';
+
+export function isValidConversationId(id: string): boolean {
+  return UUID_RE.test(id);
+}
+
+export async function verifyConversationOwnership(
+  db: D1Database,
+  conversationId: string,
+  userId: string,
+): Promise<ConversationOwnership> {
+  const row = await db.prepare(
+    'SELECT user_id FROM conversations WHERE id = ?',
+  ).bind(conversationId).first<{ user_id: string | null }>();
+
+  if (!row) return 'not_found';
+  return (row.user_id ?? 'operator') === userId ? 'owned' : 'not_owned';
+}

package/src/durable-objects/chat-session.ts

@@ -0,0 +1,251 @@
+import type { Env, MessageMetadata } from '../types.js';
+import { buildEdgeEnv } from '../edge-env.js';
+import { createIntent, dispatchStream } from '../kernel/dispatch.js';
+import type { DispatchResult, Executor } from '../kernel/types.js';
+import { isValidConversationId, verifyConversationOwnership } from './chat-session-auth.js';
+import { z } from 'zod';
+
+const ExecutorSchema = z.enum([
+  'claude',
+  'groq',
+  'workers_ai',
+  'claude_opus',
+  'gpt_oss',
+  'composite',
+]);
+
+const MessageFrameSchema = z.object({
+  type: z.literal('message'),
+  text: z.string().trim().min(1),
+  conversationId: z.string().optional(),
+  eventId: z.string().trim().min(1).optional(),
+  executor: ExecutorSchema.optional(),
+}).passthrough();
+
+type StoredMessage = {
+  id: string;
+  role: 'user' | 'assistant';
+  content: string;
+  metadata: string | null;
+  created_at: string;
+};
+
+export class ChatSession implements DurableObject {
+  constructor(
+    private readonly state: DurableObjectState,
+    private readonly env: Env,
+  ) {}
+
+  async fetch(request: Request): Promise<Response> {
+    if (request.headers.get('Upgrade')?.toLowerCase() !== 'websocket') {
+      return new Response('Expected WebSocket upgrade', { status: 426 });
+    }
+
+    const pair = new WebSocketPair();
+    const client = pair[0];
+    const server = pair[1];
+    const url = new URL(request.url);
+    const userId = this.userId();
+    const conversationId = url.searchParams.get('conversationId');
+
+    this.state.acceptWebSocket(server, ['aegis-chat']);
+    this.state.waitUntil(this.sendHistory(server, conversationId, userId));
+
+    const headers = new Headers();
+    if (request.headers.get('Sec-WebSocket-Protocol')?.split(',').map((p) => p.trim()).includes('aegis-chat')) {
+      headers.set('Sec-WebSocket-Protocol', 'aegis-chat');
+    }
+
+    return new Response(null, {
+      status: 101,
+      webSocket: client,
+      headers,
+    });
+  }
+
+  async webSocketMessage(ws: WebSocket, message: string | ArrayBuffer): Promise<void> {
+    let payload: unknown;
+    try {
+      payload = JSON.parse(typeof message === 'string' ? message : new TextDecoder().decode(message));
+    } catch {
+      this.send(ws, { type: 'error', error: 'Invalid JSON frame' });
+      return;
+    }
+
+    if (!isRecord(payload) || payload.type !== 'message') return;
+
+    const parsed = MessageFrameSchema.safeParse(payload);
+    if (!parsed.success) {
+      this.send(ws, { type: 'error', error: 'Invalid message frame' });
+      return;
+    }
+
+    const eventId = parsed.data.eventId ?? crypto.randomUUID();
+    const eventClaimed = await this.claimEvent(eventId);
+    if (!eventClaimed) {
+      this.send(ws, { type: 'error', error: 'duplicate event' });
+      return;
+    }
+
+    const text = parsed.data.text;
+    const userId = this.userId();
+    const conversationId = await this.resolveConversationId(ws, parsed.data.conversationId, userId, text);
+    if (!conversationId) return;
+
+    const userMessageId = crypto.randomUUID();
+    await this.env.DB.prepare(
+      'INSERT INTO messages (id, conversation_id, role, content) VALUES (?, ?, ?, ?)',
+    ).bind(userMessageId, conversationId, 'user', text).run();
+
+    this.send(ws, { type: 'start', conversationId });
+
+    try {
+      const edgeEnv = buildEdgeEnv(this.env);
+      const intent = createIntent(conversationId, text, {
+        forcedExecutor: parsed.data.executor as Executor | undefined,
+      });
+      const result = await dispatchStream(intent, edgeEnv, (delta) => {
+        this.send(ws, { type: 'delta', text: delta });
+      });
+
+      const assistantMessageId = crypto.randomUUID();
+      const metadata = buildMessageMetadata(result);
+
+      await this.env.DB.prepare(
+        'INSERT INTO messages (id, conversation_id, role, content, metadata) VALUES (?, ?, ?, ?, ?)',
+      ).bind(assistantMessageId, conversationId, 'assistant', result.text, JSON.stringify(metadata)).run();
+
+      await this.env.DB.prepare(
+        "UPDATE conversations SET updated_at = datetime('now') WHERE id = ?",
+      ).bind(conversationId).run();
+
+      this.send(ws, { type: 'done', conversationId, metadata: { id: assistantMessageId, ...metadata } });
+    } catch (err) {
+      const error = err instanceof Error ? err.message : String(err);
+      this.send(ws, { type: 'error', error });
+    }
+  }
+
+  private async resolveConversationId(
+    ws: WebSocket,
+    rawConversationId: unknown,
+    userId: string,
+    firstMessage: string,
+  ): Promise<string | null> {
+    if (rawConversationId !== undefined && typeof rawConversationId !== 'string') {
+      this.send(ws, { type: 'error', error: 'conversationId must be a UUID string' });
+      return null;
+    }
+
+    const conversationId = rawConversationId ?? crypto.randomUUID();
+    if (!isValidConversationId(conversationId)) {
+      this.send(ws, { type: 'error', error: 'conversationId must be a UUID string' });
+      return null;
+    }
+
+    const ownership = await verifyConversationOwnership(this.env.DB, conversationId, userId);
+    if (ownership === 'not_owned') {
+      this.send(ws, { type: 'error', error: 'conversation not found' });
+      return null;
+    }
+    if (ownership === 'owned') return conversationId;
+
+    await this.env.DB.prepare(
+      'INSERT INTO conversations (id, title, user_id) VALUES (?, ?, ?)',
+    ).bind(conversationId, firstMessage.slice(0, 100), userId).run();
+    return conversationId;
+  }
+
+  private async claimEvent(eventId: string): Promise<boolean> {
+    const existing = await this.env.DB.prepare(
+      'SELECT event_id FROM web_events WHERE event_id = ?',
+    ).bind(eventId).first();
+    if (existing) return false;
+
+    await this.env.DB.prepare(
+      'INSERT INTO web_events (event_id) VALUES (?)',
+    ).bind(eventId).run();
+    return true;
+  }
+
+  private async sendHistory(ws: WebSocket, conversationId: string | null, userId: string): Promise<void> {
+    if (!conversationId) {
+      this.send(ws, { type: 'history', conversationId: null, messages: [] });
+      return;
+    }
+    if (!isValidConversationId(conversationId)) {
+      this.send(ws, { type: 'error', error: 'conversationId must be a UUID string' });
+      return;
+    }
+
+    const ownership = await verifyConversationOwnership(this.env.DB, conversationId, userId);
+    if (ownership === 'not_owned') {
+      this.send(ws, { type: 'error', error: 'conversation not found' });
+      return;
+    }
+    if (ownership === 'not_found') {
+      this.send(ws, { type: 'history', conversationId, messages: [] });
+      return;
+    }
+
+    const rows = await this.env.DB.prepare(
+      `SELECT m.id, m.role, m.content, m.metadata, m.created_at
+       FROM messages m
+       JOIN conversations c ON c.id = m.conversation_id
+       WHERE m.conversation_id = ? AND c.user_id = ?
+       ORDER BY m.created_at ASC`,
+    ).bind(conversationId, userId).all<StoredMessage>();
+
+    this.send(ws, {
+      type: 'history',
+      conversationId,
+      messages: rows.results.map((message) => ({
+        ...message,
+        metadata: parseMetadata(message.metadata),
+      })),
+    });
+  }
+
+  private userId(): string {
+    return this.state.id.name ?? 'operator';
+  }
+
+  private send(ws: WebSocket, frame: Record<string, unknown>): void {
+    try {
+      ws.send(JSON.stringify(frame));
+    } catch {
+      // Ignore writes after the client disconnects.
+    }
+  }
+}
+
+function buildMessageMetadata(result: DispatchResult): MessageMetadata {
+  return {
+    classification: result.classification,
+    executor: result.executor,
+    procHit: result.procedureHit,
+    latencyMs: result.latency_ms,
+    cost: result.cost,
+    confidence: result.confidence,
+    reclassified: result.reclassified,
+    probeResult: result.probeResult,
+    grounded: result.grounded,
+    sources: result.sources,
+    unknowns: result.unknowns,
+    searched: result.searched,
+    unverifiedClaims: result.unverified_claims,
+  };
+}
+
+function parseMetadata(value: string | null): MessageMetadata | null {
+  if (!value) return null;
+  try {
+    return JSON.parse(value) as MessageMetadata;
+  } catch {
+    return null;
+  }
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === 'object' && value !== null;
+}