flynn/src/backends/piEmbedded.test.ts

import { mkdtempSync, rmSync, writeFileSync } from 'fs';
import { join } from 'path';
import { tmpdir } from 'os';
import { pathToFileURL } from 'url';
import { describe, expect, it } from 'vitest';

import { PiEmbeddedBackend } from './piEmbedded.js';

function createModule(source: string): { moduleUrl: string; cleanup: () => void } {
  const dir = mkdtempSync(join(tmpdir(), 'flynn-pi-embedded-'));
  const file = join(dir, 'module.mjs');
  writeFileSync(file, source, 'utf-8');
  return {
    moduleUrl: pathToFileURL(file).href,
    cleanup: () => rmSync(dir, { recursive: true, force: true }),
  };
}

describe('PiEmbeddedBackend', () => {
  it('returns text from a createAgentSession/run response object', async () => {
    const mod = createModule(`
      export function createAgentSession() {
        return {
          run(payload) {
            return { text: "pi says: " + payload.input };
          },
        };
      }
    `);

    try {
      const backend = new PiEmbeddedBackend({ module: mod.moduleUrl, timeoutMs: 2000 });
      const result = await backend.process({ prompt: 'hello', history: [] });
      expect(result).toBe('pi says: hello');
    } finally {
      mod.cleanup();
    }
  });

  it('falls back to string payload when object payload is rejected', async () => {
    const mod = createModule(`
      export function createAgentSession() {
        return {
          run(payload) {
            if (typeof payload !== "string") {
              throw new Error("expected string payload");
            }
            return "echo " + payload;
          },
        };
      }
    `);

    try {
      const backend = new PiEmbeddedBackend({ module: mod.moduleUrl, timeoutMs: 2000 });
      const result = await backend.process({ prompt: 'hello', history: [] });
      expect(result).toContain('echo USER: hello');
    } finally {
      mod.cleanup();
    }
  });

  it('injects Flynn system prompt fields into session payload in hybrid mode', async () => {
    const mod = createModule(`
      export function createAgentSession() {
        return {
          run(payload) {
            return { text: payload.systemPrompt ?? payload.system ?? "missing" };
          },
        };
      }
    `);

    try {
      const backend = new PiEmbeddedBackend({ module: mod.moduleUrl, timeoutMs: 2000, systemPromptMode: 'hybrid' });
      const result = await backend.process({
        prompt: 'hello',
        history: [],
        systemPrompt: 'SOUL + IDENTITY + USER + TOOLS',
      });
      expect(result).toBe('SOUL + IDENTITY + USER + TOOLS');
    } finally {
      mod.cleanup();
    }
  });

  it('omits Flynn system prompt injection in pi_default mode', async () => {
    const mod = createModule(`
      export function createAgentSession() {
        return {
          run(payload) {
            return { text: payload.systemPrompt ? "present" : "absent" };
          },
        };
      }
    `);

    try {
      const backend = new PiEmbeddedBackend({ module: mod.moduleUrl, timeoutMs: 2000, systemPromptMode: 'pi_default' });
      const result = await backend.process({
        prompt: 'hello',
        history: [],
        systemPrompt: 'should not be forwarded',
      });
      expect(result).toBe('absent');
    } finally {
      mod.cleanup();
    }
  });

  it('throws when module has no supported session factory', async () => {
    const mod = createModule('export const version = "0.0.0";');

    try {
      const backend = new PiEmbeddedBackend({ module: mod.moduleUrl, timeoutMs: 2000 });
      await expect(backend.process({ prompt: 'hello', history: [] }))
        .rejects.toThrow('supported runtime API');
    } finally {
      mod.cleanup();
    }
  });

  it('uses Agent class runtime when session factory exports are absent', async () => {
    const mod = createModule(`
      export class Agent {
        constructor() {
          this.state = { messages: [] };
        }
        replaceMessages(messages) {
          this.state.messages = messages.slice();
        }
        async prompt(input) {
          this.state.messages.push({ role: "user", content: [{ type: "text", text: input }] });
          this.state.messages.push({ role: "assistant", content: [{ type: "text", text: "agent says: " + input }] });
        }
      }
    `);

    try {
      const backend = new PiEmbeddedBackend({ module: mod.moduleUrl, timeoutMs: 2000 });
      const result = await backend.process({
        prompt: 'hello',
        history: [{ role: 'assistant', content: 'previous answer' }],
      });
      expect(result).toBe('agent says: hello');
    } finally {
      mod.cleanup();
    }
  });

  it('applies Flynn system prompt in Agent runtime via setSystemPrompt()', async () => {
    const mod = createModule(`
      export class Agent {
        constructor() {
          this.systemPrompt = "";
          this.state = { messages: [] };
        }
        setSystemPrompt(prompt) {
          this.systemPrompt = prompt;
        }
        async prompt(input) {
          this.state.messages.push({ role: "assistant", content: [{ type: "text", text: this.systemPrompt + " :: " + input }] });
        }
      }
    `);

    try {
      const backend = new PiEmbeddedBackend({ module: mod.moduleUrl, timeoutMs: 2000, systemPromptMode: 'flynn' });
      const result = await backend.process({
        prompt: 'hello',
        history: [],
        systemPrompt: 'use flynn prompt',
      });
      expect(result).toBe('use flynn prompt :: hello');
    } finally {
      mod.cleanup();
    }
  });

  it('surfaces agent state error when no assistant text is produced', async () => {
    const mod = createModule(`
      export class Agent {
        constructor() {
          this.state = { messages: [], error: undefined };
        }
        async prompt() {
          this.state.error = "Missing API key for default provider";
        }
      }
    `);

    try {
      const backend = new PiEmbeddedBackend({ module: mod.moduleUrl, timeoutMs: 2000 });
      await expect(backend.process({ prompt: 'hello', history: [] }))
        .rejects.toThrow('Missing API key for default provider');
    } finally {
      mod.cleanup();
    }
  });

  it('throws when module cannot be loaded', async () => {
    const backend = new PiEmbeddedBackend({ module: '/definitely/missing/pi-module.mjs', timeoutMs: 2000 });
    await expect(backend.process({ prompt: 'hello', history: [] }))
      .rejects.toThrow('Failed to load Pi embedded runtime module');
  });

  it('times out slow Pi requests', async () => {
    const mod = createModule(`
      export function createAgentSession() {
        return {
          run() {
            return new Promise(() => {});
          },
        };
      }
    `);

    try {
      const backend = new PiEmbeddedBackend({ module: mod.moduleUrl, timeoutMs: 10 });
      await expect(backend.process({ prompt: 'hello', history: [] }))
        .rejects.toThrow('timed out');
    } finally {
      mod.cleanup();
    }
  });
});