@pikku/core 0.11.2 → 0.12.1

package/dist/wirings/ai-agent/ai-agent-stream.js

@@ -0,0 +1,662 @@
+import { pikkuState, getSingletonServices } from '../../pikku-state.js';
+import { combineChannelMiddleware, wrapChannelWithMiddleware, } from '../channel/channel-middleware-runner.js';
+import { randomUUID } from 'crypto';
+import { parseWorkingMemory, deepMergeWorkingMemory, resolveMemoryServices, loadContextMessages, trimMessages, } from './ai-agent-memory.js';
+import { prepareAgentRun, resolveAgent, buildInstructions, buildToolDefs, createScopedChannel, ToolApprovalRequired, } from './ai-agent-prepare.js';
+import { createAssistantUIChannel, parseAssistantUIInput, } from './ai-agent-assistant-ui.js';
+import { resolveModelConfig } from './ai-agent-model-config.js';
+function createPersistingChannel(parent, storage, threadId) {
+    let fullText = '';
+    let stepText = '';
+    let stepToolCalls = [];
+    let stepToolResults = [];
+    const totalUsage = {
+        inputTokens: 0,
+        outputTokens: 0,
+    };
+    const flushStep = async () => {
+        if (!storage)
+            return;
+        const text = stepText;
+        const calls = stepToolCalls;
+        const results = stepToolResults;
+        stepText = '';
+        stepToolCalls = [];
+        stepToolResults = [];
+        if (text || calls.length > 0) {
+            const assistantMsg = {
+                id: randomUUID(),
+                role: 'assistant',
+                content: text || undefined,
+                toolCalls: calls.length > 0 ? calls : undefined,
+                createdAt: new Date(),
+            };
+            const messages = [assistantMsg];
+            if (results.length > 0) {
+                messages.push({
+                    id: randomUUID(),
+                    role: 'tool',
+                    toolResults: results,
+                    createdAt: new Date(),
+                });
+            }
+            await storage.saveMessages(threadId, messages);
+        }
+    };
+    const channel = {
+        channelId: parent.channelId,
+        openingData: parent.openingData,
+        get state() {
+            return parent.state;
+        },
+        get fullText() {
+            return fullText;
+        },
+        get totalUsage() {
+            return totalUsage;
+        },
+        flush: flushStep,
+        close: () => parent.close(),
+        send: (event) => {
+            if (storage) {
+                switch (event.type) {
+                    case 'text-delta':
+                        stepText += event.text;
+                        fullText += event.text;
+                        break;
+                    case 'tool-call':
+                        stepToolCalls.push({
+                            id: event.toolCallId,
+                            name: event.toolName,
+                            args: event.args,
+                        });
+                        break;
+                    case 'tool-result':
+                        stepToolResults.push({
+                            id: event.toolCallId,
+                            name: event.toolName,
+                            result: typeof event.result === 'string'
+                                ? event.result
+                                : JSON.stringify(event.result),
+                        });
+                        break;
+                    case 'usage':
+                        totalUsage.inputTokens += event.tokens.input;
+                        totalUsage.outputTokens += event.tokens.output;
+                        if (event.model)
+                            totalUsage.model = event.model;
+                        flushStep();
+                        break;
+                    case 'done':
+                        flushStep();
+                        break;
+                }
+            }
+            parent.send(event);
+        },
+    };
+    return channel;
+}
+async function postStreamCleanup(persistingChannel, aiMiddlewares, singletonServices, storage, memoryConfig, threadId, workingMemorySchemaName, messages, aiRunState, runId) {
+    const usage = persistingChannel.totalUsage;
+    let outputText = persistingChannel.fullText;
+    let outputMessages = messages;
+    for (let i = aiMiddlewares.length - 1; i >= 0; i--) {
+        const mw = aiMiddlewares[i];
+        if (mw.modifyOutput) {
+            const result = await mw.modifyOutput(singletonServices, {
+                text: outputText,
+                messages: outputMessages,
+                usage: {
+                    inputTokens: usage.inputTokens,
+                    outputTokens: usage.outputTokens,
+                },
+            });
+            outputText = result.text;
+            outputMessages = result.messages;
+        }
+    }
+    if (storage && memoryConfig?.workingMemory && outputText) {
+        const parsed = parseWorkingMemory(outputText);
+        if (parsed) {
+            const existing = (await storage.getWorkingMemory(threadId, 'thread')) ?? {};
+            const merged = deepMergeWorkingMemory(existing, parsed);
+            if (singletonServices.schema && workingMemorySchemaName) {
+                try {
+                    await singletonServices.schema.validateSchema(workingMemorySchemaName, merged);
+                }
+                catch (err) {
+                    singletonServices.logger.warn(`Working memory validation failed: ${err instanceof Error ? err.message : String(err)}`);
+                }
+            }
+            await storage.saveWorkingMemory(threadId, 'thread', merged);
+        }
+    }
+    await aiRunState.updateRun(runId, {
+        status: 'completed',
+        ...(usage.model
+            ? {
+                usage: {
+                    inputTokens: usage.inputTokens,
+                    outputTokens: usage.outputTokens,
+                    model: usage.model,
+                },
+            }
+            : {}),
+    });
+}
+async function runStreamStepLoop(params) {
+    const { agent, runnerParams, maxSteps, agentRunner, persistingChannel, channel, runId, aiMiddlewares, } = params;
+    const singletonServices = getSingletonServices();
+    for (let step = 0; step < maxSteps; step++) {
+        if (agent.prepareStep) {
+            let stopped = false;
+            await agent.prepareStep({
+                stepNumber: step,
+                messages: runnerParams.messages,
+                tools: runnerParams.tools,
+                toolChoice: runnerParams.toolChoice,
+                model: runnerParams.model,
+                stop: () => {
+                    stopped = true;
+                },
+            });
+            if (stopped)
+                break;
+        }
+        channel.send({ type: 'step-start', stepNumber: step });
+        const stepResult = await agentRunner.stream(runnerParams, persistingChannel);
+        for (const mw of aiMiddlewares) {
+            if (mw.afterStep) {
+                await mw.afterStep(singletonServices, {
+                    stepNumber: step,
+                    text: stepResult.text,
+                    toolCalls: stepResult.toolCalls,
+                    toolResults: stepResult.toolResults,
+                    usage: stepResult.usage,
+                    finishReason: stepResult.finishReason,
+                });
+            }
+        }
+        if (stepResult.toolCalls.length === 0)
+            break;
+        const approvalNeeded = checkForApproval(stepResult, runnerParams.tools, runId);
+        if (approvalNeeded) {
+            return { outcome: 'approval', approval: approvalNeeded };
+        }
+        appendStepMessages(runnerParams, stepResult);
+    }
+    return { outcome: 'done' };
+}
+function checkForApproval(stepResult, tools, runId) {
+    for (const tc of stepResult.toolCalls) {
+        const toolDef = tools.find((t) => t.name === tc.toolName);
+        if (toolDef?.needsApproval) {
+            return new ToolApprovalRequired(tc.toolCallId, tc.toolName, tc.args);
+        }
+        const tr = stepResult.toolResults.find((r) => r.toolCallId === tc.toolCallId);
+        if (tr?.result &&
+            typeof tr.result === 'object' &&
+            '__approvalRequired' in tr.result) {
+            const r = tr.result;
+            return new ToolApprovalRequired(tc.toolCallId, r.toolName, r.args, r.reason, r.displayToolName, r.displayArgs, r.agentRunId);
+        }
+    }
+    return null;
+}
+function appendStepMessages(runnerParams, stepResult) {
+    const assistantMsg = {
+        id: randomUUID(),
+        role: 'assistant',
+        content: stepResult.text || undefined,
+        toolCalls: stepResult.toolCalls.length > 0
+            ? stepResult.toolCalls.map((tc) => ({
+                id: tc.toolCallId,
+                name: tc.toolName,
+                args: tc.args,
+            }))
+            : undefined,
+        createdAt: new Date(),
+    };
+    runnerParams.messages.push(assistantMsg);
+    if (stepResult.toolResults.length > 0) {
+        const toolMsg = {
+            id: randomUUID(),
+            role: 'tool',
+            toolResults: stepResult.toolResults.map((tr) => ({
+                id: tr.toolCallId,
+                name: tr.toolName,
+                result: typeof tr.result === 'string' ? tr.result : JSON.stringify(tr.result),
+            })),
+            createdAt: new Date(),
+        };
+        runnerParams.messages.push(toolMsg);
+    }
+}
+function handleApproval(err, runId, channel, aiRunState, persistingChannel) {
+    return (async () => {
+        await persistingChannel.flush();
+        const pendingApproval = err.agentRunId
+            ? {
+                type: 'agent-call',
+                toolCallId: err.toolCallId,
+                agentName: err.toolName,
+                agentRunId: err.agentRunId,
+                displayToolName: err.displayToolName ?? err.toolName,
+                displayArgs: err.displayArgs ?? err.args,
+            }
+            : {
+                type: 'tool-call',
+                toolCallId: err.toolCallId,
+                toolName: err.toolName,
+                args: err.args,
+            };
+        await aiRunState.updateRun(runId, {
+            status: 'suspended',
+            suspendReason: 'approval',
+            pendingApprovals: [pendingApproval],
+        });
+        const approvalEvent = {
+            type: 'approval-request',
+            toolCallId: err.toolCallId,
+            toolName: err.displayToolName ?? err.toolName,
+            args: err.displayArgs ?? err.args,
+            reason: err.reason,
+            runId,
+        };
+        channel.send(approvalEvent);
+        channel.send({ type: 'done' });
+        channel.close();
+    })();
+}
+export async function streamAIAgent(agentName, input, channel, params, agentSessionMap, options) {
+    const sessionMap = agentSessionMap ?? new Map();
+    const { agent: resolvedAgentForProtocol } = resolveAgent(agentName);
+    const inputRecord = input;
+    const useUIMessageStream = resolvedAgentForProtocol.protocol === 'ui-message-stream' ||
+        Array.isArray(inputRecord.messages);
+    let normalizedInput;
+    if (useUIMessageStream) {
+        normalizedInput = parseAssistantUIInput(inputRecord);
+        channel = createAssistantUIChannel(channel);
+    }
+    else {
+        normalizedInput = input;
+    }
+    const streamContext = { channel, options };
+    const { agent, packageName, resolvedName, agentRunner, storage, memoryConfig, threadId, userMessage, runnerParams, maxSteps, missingRpcs, workingMemorySchemaName, } = await prepareAgentRun(agentName, normalizedInput, params, sessionMap, streamContext);
+    const singletonServices = getSingletonServices();
+    const { aiRunState } = singletonServices;
+    if (!aiRunState) {
+        throw new Error('AIRunStateService not available in singletonServices');
+    }
+    if (missingRpcs.length > 0) {
+        await aiRunState.createRun({
+            agentName,
+            threadId,
+            resourceId: normalizedInput.resourceId,
+            status: 'suspended',
+            suspendReason: 'rpc-missing',
+            missingRpcs,
+            usage: { inputTokens: 0, outputTokens: 0, model: agent.model },
+            createdAt: new Date(),
+            updatedAt: new Date(),
+        });
+        channel.send({ type: 'suspended', reason: 'rpc-missing', missingRpcs });
+        channel.send({ type: 'done' });
+        return;
+    }
+    const aiMiddlewares = agent.aiMiddleware ?? [];
+    let modifiedMessages = runnerParams.messages;
+    let modifiedInstructions = runnerParams.instructions;
+    for (const mw of aiMiddlewares) {
+        if (mw.modifyInput) {
+            const result = await mw.modifyInput(singletonServices, {
+                messages: modifiedMessages,
+                instructions: modifiedInstructions,
+            });
+            modifiedMessages = result.messages;
+            modifiedInstructions = result.instructions;
+        }
+    }
+    runnerParams.messages = modifiedMessages;
+    runnerParams.instructions = modifiedInstructions;
+    const runId = await aiRunState.createRun({
+        agentName,
+        threadId,
+        resourceId: normalizedInput.resourceId,
+        status: 'running',
+        usage: { inputTokens: 0, outputTokens: 0, model: agent.model },
+        createdAt: new Date(),
+        updatedAt: new Date(),
+    });
+    if (storage) {
+        await storage.saveMessages(threadId, [userMessage]);
+    }
+    const streamMiddleware = aiMiddlewares
+        .filter((mw) => mw.modifyOutputStream)
+        .map((mw) => {
+        const state = {};
+        const allEvents = [];
+        return async (services, event, next) => {
+            allEvents.push(event);
+            const result = await mw.modifyOutputStream(services, {
+                event,
+                allEvents,
+                state,
+            });
+            if (result != null)
+                await next(result);
+        };
+    });
+    const agentsMeta = pikkuState(packageName, 'agent', 'agentsMeta');
+    const meta = agentsMeta[resolvedName];
+    const allChannelMiddleware = combineChannelMiddleware('agent', `stream:${agentName}`, {
+        wireInheritedChannelMiddleware: meta?.channelMiddleware,
+        wireChannelMiddleware: [
+            ...(agent.channelMiddleware ?? []),
+            ...streamMiddleware,
+        ],
+    });
+    const wrappedChannel = allChannelMiddleware.length > 0
+        ? wrapChannelWithMiddleware({ channel }, singletonServices, allChannelMiddleware).channel
+        : channel;
+    const persistingChannel = createPersistingChannel(wrappedChannel, storage, threadId);
+    try {
+        const loopResult = await runStreamStepLoop({
+            agent,
+            runnerParams,
+            maxSteps,
+            agentRunner,
+            persistingChannel,
+            channel,
+            runId,
+            aiMiddlewares,
+        });
+        if (loopResult.outcome === 'approval') {
+            await handleApproval(loopResult.approval, runId, channel, aiRunState, persistingChannel);
+            return;
+        }
+        await postStreamCleanup(persistingChannel, aiMiddlewares, singletonServices, storage, memoryConfig, threadId, workingMemorySchemaName, runnerParams.messages, aiRunState, runId);
+        channel.send({ type: 'done' });
+        channel.close();
+    }
+    catch (err) {
+        for (const mw of aiMiddlewares) {
+            if (mw.onError) {
+                try {
+                    await mw.onError(singletonServices, {
+                        error: err instanceof Error ? err : new Error(String(err)),
+                        stepNumber: -1,
+                        messages: runnerParams.messages,
+                    });
+                }
+                catch {
+                    // onError hooks must not affect error flow
+                }
+            }
+        }
+        await aiRunState.updateRun(runId, { status: 'failed' });
+        channel.send({
+            type: 'error',
+            message: err instanceof Error ? err.message : String(err),
+        });
+        channel.send({ type: 'done' });
+        channel.close();
+    }
+}
+export async function resumeAIAgent(input, channel, params, options) {
+    const singletonServices = getSingletonServices();
+    const { aiRunState } = singletonServices;
+    if (!aiRunState) {
+        throw new Error('AIRunStateService not available in singletonServices');
+    }
+    const run = await aiRunState.getRun(input.runId);
+    if (!run) {
+        throw new Error(`No run found for runId ${input.runId}`);
+    }
+    const pending = run.pendingApprovals?.find((p) => p.toolCallId === input.toolCallId);
+    if (!pending) {
+        throw new Error(`No pending approval for toolCallId ${input.toolCallId} on run ${input.runId}`);
+    }
+    await aiRunState.resolveApproval(input.toolCallId, input.approved ? 'approved' : 'denied');
+    const { agent, packageName, resolvedName } = resolveAgent(run.agentName);
+    const { storage } = resolveMemoryServices(agent, singletonServices);
+    const memoryConfig = agent.memory;
+    const agentRunner = singletonServices.aiAgentRunner;
+    if (!agentRunner) {
+        throw new Error('AIAgentRunnerService not available in singletonServices');
+    }
+    if (!input.approved) {
+        if (pending.type === 'agent-call') {
+            await aiRunState.updateRun(pending.agentRunId, { status: 'failed' });
+        }
+        const denialResult = 'Tool call was not approved by the user';
+        if (storage) {
+            await storage.saveMessages(run.threadId, [
+                {
+                    id: randomUUID(),
+                    role: 'tool',
+                    toolResults: [
+                        {
+                            id: input.toolCallId,
+                            name: pending.type === 'tool-call'
+                                ? pending.toolName
+                                : pending.agentName,
+                            result: denialResult,
+                        },
+                    ],
+                    createdAt: new Date(),
+                },
+            ]);
+        }
+        await aiRunState.updateRun(run.runId, { status: 'running' });
+        await continueAfterToolResult(run, agent, packageName, resolvedName, storage, memoryConfig, agentRunner, channel, params, aiRunState, options);
+        return;
+    }
+    if (pending.type === 'agent-call') {
+        const subRun = await aiRunState.getRun(pending.agentRunId);
+        if (!subRun) {
+            throw new Error(`Sub-agent run not found: ${pending.agentRunId}`);
+        }
+        const subPending = subRun.pendingApprovals?.[0];
+        if (!subPending) {
+            throw new Error(`No pending approval on sub-agent run ${pending.agentRunId}`);
+        }
+        const subChannel = createScopedChannel(channel, subRun.agentName, 'resume');
+        channel.send({
+            type: 'agent-call',
+            agentName: subRun.agentName,
+            session: 'resume',
+            input: null,
+        });
+        await resumeAIAgent({
+            runId: pending.agentRunId,
+            toolCallId: subPending.toolCallId,
+            approved: true,
+        }, subChannel, params, options);
+        channel.send({
+            type: 'agent-result',
+            agentName: subRun.agentName,
+            session: 'resume',
+            result: null,
+        });
+        if (storage) {
+            await storage.saveMessages(run.threadId, [
+                {
+                    id: randomUUID(),
+                    role: 'tool',
+                    toolResults: [
+                        {
+                            id: input.toolCallId,
+                            name: pending.agentName,
+                            result: 'Sub-agent completed successfully',
+                        },
+                    ],
+                    createdAt: new Date(),
+                },
+            ]);
+        }
+    }
+    else {
+        const streamContext = {
+            channel,
+            options: { ...options, requiresToolApproval: false },
+        };
+        const aiMiddlewaresForResume = agent.aiMiddleware ?? [];
+        const { tools } = buildToolDefs(params, new Map(), run.resourceId, resolvedName, packageName, streamContext, aiMiddlewaresForResume);
+        const matchingTool = tools.find((t) => t.name === pending.toolName);
+        if (!matchingTool) {
+            throw new Error(`Tool "${pending.toolName}" not found in agent definition`);
+        }
+        const toolArgs = typeof pending.args === 'string' ? JSON.parse(pending.args) : pending.args;
+        let toolResult;
+        try {
+            toolResult = await matchingTool.execute(toolArgs);
+        }
+        catch (execErr) {
+            toolResult = `Error: ${execErr instanceof Error ? execErr.message : String(execErr)}`;
+        }
+        const resultStr = typeof toolResult === 'string' ? toolResult : JSON.stringify(toolResult);
+        if (storage) {
+            await storage.saveMessages(run.threadId, [
+                {
+                    id: randomUUID(),
+                    role: 'tool',
+                    toolResults: [
+                        {
+                            id: input.toolCallId,
+                            name: pending.toolName,
+                            result: resultStr,
+                        },
+                    ],
+                    createdAt: new Date(),
+                },
+            ]);
+        }
+        channel.send({
+            type: 'tool-result',
+            toolCallId: input.toolCallId,
+            toolName: pending.toolName,
+            result: toolResult,
+        });
+    }
+    await aiRunState.updateRun(run.runId, { status: 'running' });
+    await continueAfterToolResult(run, agent, packageName, resolvedName, storage, memoryConfig, agentRunner, channel, params, aiRunState, options);
+}
+async function continueAfterToolResult(run, agent, packageName, resolvedName, storage, memoryConfig, agentRunner, channel, params, aiRunState, options) {
+    const singletonServices = getSingletonServices();
+    const agentsMeta = pikkuState(packageName, 'agent', 'agentsMeta');
+    const meta = agentsMeta[resolvedName];
+    const workingMemorySchemaName = meta?.workingMemorySchema ?? null;
+    const messages = storage
+        ? await storage.getMessages(run.threadId, {
+            lastN: memoryConfig?.lastMessages ?? 20,
+        })
+        : [];
+    const workingMemoryJsonSchema = workingMemorySchemaName
+        ? pikkuState(packageName, 'misc', 'schemas').get(workingMemorySchemaName)
+        : undefined;
+    const contextMessages = await loadContextMessages(memoryConfig, storage, { message: '', threadId: run.threadId, resourceId: run.resourceId }, workingMemoryJsonSchema);
+    const allMessages = [...contextMessages, ...messages];
+    const trimmedMessages = trimMessages(allMessages);
+    const instructions = buildInstructions(resolvedName, packageName);
+    const aiMiddlewares = agent.aiMiddleware ?? [];
+    let modifiedMessages = trimmedMessages;
+    let modifiedInstructions = instructions;
+    for (const mw of aiMiddlewares) {
+        if (mw.modifyInput) {
+            const result = await mw.modifyInput(singletonServices, {
+                messages: modifiedMessages,
+                instructions: modifiedInstructions,
+            });
+            modifiedMessages = result.messages;
+            modifiedInstructions = result.instructions;
+        }
+    }
+    const streamMiddleware = aiMiddlewares
+        .filter((mw) => mw.modifyOutputStream)
+        .map((mw) => {
+        const state = {};
+        const allEvents = [];
+        return async (services, event, next) => {
+            allEvents.push(event);
+            const result = await mw.modifyOutputStream(services, {
+                event,
+                allEvents,
+                state,
+            });
+            if (result != null)
+                await next(result);
+        };
+    });
+    const allChannelMiddleware = combineChannelMiddleware('agent', `stream:${run.agentName}`, {
+        wireInheritedChannelMiddleware: meta?.channelMiddleware,
+        wireChannelMiddleware: [
+            ...(agent.channelMiddleware ?? []),
+            ...streamMiddleware,
+        ],
+    });
+    const wrappedChannel = allChannelMiddleware.length > 0
+        ? wrapChannelWithMiddleware({ channel }, singletonServices, allChannelMiddleware).channel
+        : channel;
+    const persistingChannel = createPersistingChannel(wrappedChannel, storage, run.threadId);
+    const streamContext = { channel, options };
+    const resumeTools = buildToolDefs(params, new Map(), run.resourceId, resolvedName, packageName, streamContext, aiMiddlewares).tools;
+    const resolved = resolveModelConfig(resolvedName, agent);
+    const maxSteps = resolved.maxSteps ?? 10;
+    const runnerParams = {
+        model: resolved.model,
+        temperature: resolved.temperature,
+        instructions: modifiedInstructions,
+        messages: modifiedMessages,
+        tools: resumeTools,
+        maxSteps: 1,
+        toolChoice: (agent.toolChoice ?? 'auto'),
+        outputSchema: meta?.outputSchema
+            ? pikkuState(packageName, 'misc', 'schemas').get(meta.outputSchema)
+            : undefined,
+    };
+    try {
+        const loopResult = await runStreamStepLoop({
+            agent,
+            runnerParams,
+            maxSteps,
+            agentRunner,
+            persistingChannel,
+            channel,
+            runId: run.runId,
+            aiMiddlewares,
+        });
+        if (loopResult.outcome === 'approval') {
+            await handleApproval(loopResult.approval, run.runId, channel, aiRunState, persistingChannel);
+            return;
+        }
+        await postStreamCleanup(persistingChannel, aiMiddlewares, singletonServices, storage, memoryConfig, run.threadId, workingMemorySchemaName, runnerParams.messages, aiRunState, run.runId);
+        channel.send({ type: 'done' });
+        channel.close();
+    }
+    catch (err) {
+        for (const mw of aiMiddlewares) {
+            if (mw.onError) {
+                try {
+                    await mw.onError(singletonServices, {
+                        error: err instanceof Error ? err : new Error(String(err)),
+                        stepNumber: -1,
+                        messages: runnerParams.messages,
+                    });
+                }
+                catch {
+                    // onError hooks must not affect error flow
+                }
+            }
+        }
+        await aiRunState.updateRun(run.runId, { status: 'failed' });
+        channel.send({
+            type: 'error',
+            message: err instanceof Error ? err.message : String(err),
+        });
+        channel.send({ type: 'done' });
+        channel.close();
+    }
+}