Integrate Atlas Chat

/**
 * AtlasClient — stateless HTTP/SSE transport layer.
 *
 * Single responsibility: send chat payloads to the Cognite AI agent API
 * and parse the response (JSON or Server-Sent Events).
 */

import type { CogniteClient } from '@cognite/sdk';
import type { Agent, ChatPayload, RawAgentResponse, StreamCallbacks } from './types';

const CDF_API_VERSION = 'alpha';
const AGENTS_API_VERSION = 'beta';

export class AtlasClient {
  private readonly client: CogniteClient;

  constructor(client: CogniteClient) {
    this.client = client;
  }

  async listAgents(): Promise<Agent[]> {
    const { data } = await this.client.get<{ items: Agent[] }>(
      `/api/v1/projects/${this.client.project}/ai/agents`,
      { headers: { 'cdf-version': AGENTS_API_VERSION } },
    );
    return data.items;
  }

  async getAgentById(externalId: string): Promise<Agent | null> {
    const { data } = await this.client.post<{ items: Agent[] }>(
      `/api/v1/projects/${this.client.project}/ai/agents/byids`,
      { data: { items: [{ externalId }] }, headers: { 'cdf-version': AGENTS_API_VERSION } },
    );
    return data.items[0] ?? null;
  }

  /**
   * Post a chat payload and parse the response (JSON or SSE).
   * @param agentExternalId — used as a fallback identifier when the SSE result event omits agent IDs.
   */
  async post(
    payload: ChatPayload,
    agentExternalId: string,
    callbacks?: StreamCallbacks,
    signal?: AbortSignal,
  ): Promise<RawAgentResponse> {
    const url = `${this.client.getBaseUrl()}/api/v1/projects/${this.client.project}/ai/internal/agents/chat`;

    const response = await fetch(url, {
      method: 'POST',
      headers: {
        ...this.client.getDefaultRequestHeaders(),
        'Content-Type': 'application/json',
        'cdf-version': CDF_API_VERSION,
      },
      body: JSON.stringify(payload),
      signal,
    });

    if (!response.ok) {
      const errorData = await response.json().catch(() => ({}));
      throw new Error(
        `Agent chat API error: ${response.status} - ${JSON.stringify(errorData)}`,
      );
    }

    const contentType = response.headers.get('content-type') || '';

    if (
      contentType.includes('text/event-stream') ||
      contentType.includes('text/plain')
    ) {
      return this.parseSSE(response, agentExternalId, callbacks);
    }

    return await response.json();
  }

  /**
   * Parse a Server-Sent Events streaming response.
   */
  private async parseSSE(
    response: Response,
    agentExternalId: string,
    callbacks?: StreamCallbacks,
  ): Promise<RawAgentResponse> {
    const reader = response.body?.getReader();
    if (!reader) {
      throw new Error('Response body is not readable');
    }

    const decoder = new TextDecoder();
    let buffer = '';
    let finalResponse: RawAgentResponse | null = null;

    const processLine = (line: string) => {
      if (!line.startsWith('data: ')) return;

      const dataStr = line.slice(6).trim();
      if (dataStr === '[DONE]') return '[DONE]' as const;

      try {
        const data = JSON.parse(dataStr);
        const sseResponse = data.response;

        if (!sseResponse) return;

        if (sseResponse.type === 'progress' && sseResponse.content) {
          callbacks?.onProgress?.(sseResponse.content);
        } else if (
          sseResponse.type === 'responseChunk' &&
          sseResponse.content
        ) {
          callbacks?.onChunk?.(sseResponse.content);
        } else if (sseResponse.type === 'result') {
          finalResponse = {
            agentId: data.agentId || agentExternalId,
            agentExternalId: data.agentExternalId || agentExternalId,
            response: {
              type: 'result',
              cursor: sseResponse.cursor,
              messages: sseResponse.messages || [],
            },
          };
        }
      } catch {
        // Skip unparseable SSE lines
      }
    };

    try {
      while (true) {
        const { done, value } = await reader.read();

        if (done) break;

        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split('\n');
        buffer = lines.pop() || '';

        if (lines.some((line) => processLine(line) === '[DONE]')) break;
      }
    } finally {
      reader.releaseLock();
    }

    if (!finalResponse) {
      throw new Error('No result response received from streaming API');
    }

    return finalResponse;
  }
}

/**
 * Atlas Agent Client — public API.
 *
 * Core library for communicating with Cognite Atlas AI agents.
 * Self-contained, zero imports from outside this directory (except external packages).
 *
 * React hook is a separate import path:
 *   import { useAtlasChat } from '@cognite/dune-utils/atlas-agent/react';
 */

// Core
export { AtlasSession } from './session';
export { AtlasClient } from './client';

// TypeBox re-exports for convenience
export { Type } from '@sinclair/typebox';
export type { Static, TSchema } from '@sinclair/typebox';

// Types
export type {
  Agent,
  AgentToolConfig,
  AtlasTool,
  AtlasToolResult,
  AtlasResponse,
  AtlasSessionConfig,
  ToolCall,
  StreamCallbacks,
  ApiToolDefinition,
  PythonRuntime,
} from './types';

/**
 * usePyodideRuntime — React hook for managing PyodideRuntime lifecycle.
 *
 * Separate entry point so the core atlas-agent bundle stays Pyodide-free.
 */

import { useState, useRef, useEffect } from 'react';
import type { CogniteClient } from '@cognite/sdk';
import { getGlobalPyodideRuntime } from './pyodide-runtime';
import type { PyodideRuntimeConfig } from './pyodide-runtime';
import type { PythonRuntime } from './types';

// ============================================================================
// Types
// ============================================================================

export interface PyodideProgress {
  stage: string;
  percent: number;
}

export interface UsePyodideRuntimeOptions {
  /** The `loadPyodide` function from the `pyodide` package. */
  loadPyodide: PyodideRuntimeConfig['loadPyodide'];
  /** CogniteClient for SDK credential injection. `null` disables initialization. */
  client: CogniteClient | null;
  /** Additional Python packages to install via micropip. */
  requirements?: string[];
  /** CDN URL for Pyodide files. */
  cdnUrl?: string;
}

export interface UsePyodideRuntimeReturn {
  /** The initialized runtime, or undefined if not yet ready. */
  runtime: PythonRuntime | undefined;
  /** True while Pyodide is loading / initializing. */
  loading: boolean;
  /** Error message if initialization failed. */
  error: string | null;
  /** Current initialization progress. */
  progress: PyodideProgress;
  /** Convenience: true when runtime is ready to use. */
  isReady: boolean;
}

// ============================================================================
// Hook
// ============================================================================

const DEFAULT_BASE_URL = 'https://api.cognitedata.com';

/**
 * Manages PyodideRuntime initialization lifecycle.
 *
 * Loads Pyodide, installs packages, injects Cognite SDK credentials,
 * and returns a ready-to-use `PythonRuntime` with loading/error state.
 *
 * ```tsx
 * import { loadPyodide } from 'pyodide';
 *
 * const { runtime, loading, progress, isReady } = usePyodideRuntime({
 *   loadPyodide,
 *   client: sdk,
 *   requirements: ['pandas', 'numpy'],
 * });
 * ```
 */
export function usePyodideRuntime(
  options: UsePyodideRuntimeOptions,
): UsePyodideRuntimeReturn {
  const { client } = options;

  const [runtime, setRuntime] = useState<PythonRuntime>();
  const [loading, setLoading] = useState(true);
  const [error, setError] = useState<string | null>(null);
  const [progress, setProgress] = useState<PyodideProgress>({ stage: '', percent: 0 });

  // Refs for values that shouldn't trigger re-initialization
  const loadPyodideRef = useRef(options.loadPyodide);
  const requirementsRef = useRef(options.requirements);
  const cdnUrlRef = useRef(options.cdnUrl);
  loadPyodideRef.current = options.loadPyodide;
  requirementsRef.current = options.requirements;
  cdnUrlRef.current = options.cdnUrl;

  useEffect(() => {
    if (!client) {
      setLoading(false);
      return;
    }

    let mounted = true;

    (async () => {
      try {
        setLoading(true);
        setError(null);

        const instance = getGlobalPyodideRuntime({
          loadPyodide: loadPyodideRef.current,
          requirements: requirementsRef.current,
          cdnUrl: cdnUrlRef.current,
          onProgress: (stage, percent) => {
            if (mounted) setProgress({ stage, percent });
          },
        });

        if (!instance.isInitialized) {
          if (mounted) setProgress({ stage: 'Initializing...', percent: 0 });

          const headers = client.getDefaultRequestHeaders();
          const token = headers.Authorization?.split(' ')[1] ?? '';

          await instance.initialize({
            project: client.project,
            baseUrl: client.getBaseUrl?.() ?? DEFAULT_BASE_URL,
            token,
          });
        }

        if (mounted) {
          setRuntime(instance);
          setProgress({ stage: 'Ready', percent: 100 });
        }
      } catch (err) {
        if (mounted) {
          setError(err instanceof Error ? err.message : String(err));
        }
      } finally {
        if (mounted) setLoading(false);
      }
    })();

    return () => {
      mounted = false;
    };
  }, [client]);

  return {
    runtime,
    loading,
    error,
    progress,
    isReady: !loading && !error && runtime !== undefined,
  };
}

/**
 * PyodideRuntime — browser-based Python execution via Pyodide.
 *
 * Wraps Pyodide loading, package installation, Cognite SDK setup,
 * and localStorage caching into a clean PythonRuntime implementation.
 *
 * The consumer owns the 'pyodide' npm package — they pass `loadPyodide`
 * as a config parameter so this module has no hard dependency on it.
 */

import type { PythonRuntime } from './types';

// ============================================================================
// Minimal Pyodide Interfaces (avoids 'pyodide' package dependency)
// ============================================================================

interface PyodideGlobals {
  get(name: string): unknown;
  set(name: string, value: unknown): void;
}

/** Subset of PyodideInterface that this module uses. */
export interface PyodideInstance {
  loadPackage(packages: string[]): Promise<void>;
  runPython(code: string): unknown;
  runPythonAsync(code: string): Promise<unknown>;
  globals: PyodideGlobals;
  pyimport(name: string): unknown;
}

interface Micropip {
  install(packages: string | string[]): Promise<void>;
}

// ============================================================================
// Constants
// ============================================================================

const DEFAULT_CDN_URL = 'https://cdn.jsdelivr.net/pyodide/v0.29.3/full/';
const CACHE_KEY = 'dune_pyodide_initialized';
const CACHE_VERSION = 'v1';

// ============================================================================
// Config Types
// ============================================================================

export interface PyodideRuntimeConfig {
  /** The `loadPyodide` function from the `pyodide` package. */
  loadPyodide: (options: { indexURL: string }) => Promise<PyodideInstance>;
  /** CDN URL for Pyodide files. Defaults to jsDelivr v0.29.3. */
  cdnUrl?: string;
  /** Additional Python packages to install via micropip. */
  requirements?: string[];
  /** Progress callback for initialization stages. */
  onProgress?: (stage: string, percent: number) => void;
}

export interface PyodideSDKConfig {
  project: string;
  baseUrl: string;
  token: string;
}

// ============================================================================
// Python Utility Code (injected at init)
// ============================================================================

const PYTHON_UTILS = `
import json

def _serialize_cognite_object(obj, depth=0):
    if depth > 10: return str(obj)
    for attr in ('dump', 'as_dict'):
        fn = getattr(obj, attr, None)
        if fn:
            try: return fn()
            except: pass
    if isinstance(obj, dict):
        return {k: _serialize_cognite_object(v, depth+1) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_serialize_cognite_object(i, depth+1) for i in obj]
    if isinstance(obj, (str, int, float, bool, type(None))):
        return obj
    d = getattr(obj, '__dict__', None)
    if d is not None:
        try: return _serialize_cognite_object(d, depth+1)
        except: pass
    return str(obj)

def as_json_string(value):
    return json.dumps(_serialize_cognite_object(value))
`;

// ============================================================================
// Cache Helpers
// ============================================================================

function isCacheValid(): boolean {
  try {
    return localStorage.getItem(CACHE_KEY) === CACHE_VERSION;
  } catch {
    return false;
  }
}

function markCacheValid(): void {
  try {
    localStorage.setItem(CACHE_KEY, CACHE_VERSION);
  } catch {
    /* localStorage unavailable */
  }
}

/** Clear the Pyodide package cache — forces re-download on next init. */
export function clearPyodideCache(): void {
  try {
    localStorage.removeItem(CACHE_KEY);
  } catch {
    /* localStorage unavailable */
  }
}

// ============================================================================
// PyProxy Detection (structural — avoids importing from 'pyodide')
// ============================================================================

function isPyProxy(value: unknown): value is { destroy(): void } {
  return (
    value != null &&
    typeof value === 'object' &&
    'destroy' in value &&
    typeof (value as Record<string, unknown>).destroy === 'function'
  );
}

function destroyIfPyProxy(value: unknown): void {
  if (isPyProxy(value)) {
    value.destroy();
  }
}

// ============================================================================
// PyodideRuntime Class
// ============================================================================

/**
 * PythonRuntime backed by Pyodide. Handles loading, package installation,
 * Cognite SDK credential injection, caching, and PyProxy conversion.
 */
export class PyodideRuntime implements PythonRuntime {
  private pyodide?: PyodideInstance;
  private micropip?: Micropip;
  private _initialized = false;
  private readonly config: PyodideRuntimeConfig;

  constructor(config: PyodideRuntimeConfig) {
    this.config = config;
  }

  get isInitialized(): boolean {
    return this._initialized;
  }

  /**
   * Load Pyodide, install packages, and set up the Cognite SDK.
   * Safe to call multiple times — subsequent calls are no-ops.
   */
  async initialize(sdk: PyodideSDKConfig): Promise<void> {
    if (this._initialized) return;

    const report = this.config.onProgress ?? (() => {});
    const cdnUrl = this.config.cdnUrl ?? DEFAULT_CDN_URL;

    // 1. Load Pyodide
    report('Loading Pyodide...', 10);
    this.pyodide = await this.config.loadPyodide({ indexURL: cdnUrl });
    report('Pyodide loaded', 30);

    // 2. Core packages (micropip + HTTP patching)
    report('Loading core packages...', 40);
    await this.pyodide.loadPackage(['micropip', 'pyodide-http']);
    await this.pyodide.runPythonAsync(`
try:
    import pyodide_http
    pyodide_http.patch_all()
except Exception:
    pass
`);
    this.micropip = this.pyodide.pyimport('micropip') as Micropip;

    // 3. Cognite SDK
    const verb = isCacheValid() ? 'Loading' : 'Downloading';
    report(`${verb} cognite-sdk...`, 60);
    await this.micropip.install('cognite-sdk');
    if (!isCacheValid()) markCacheValid();
    report('cognite-sdk ready', 80);

    // 4. Additional packages
    const reqs = this.config.requirements ?? [];
    if (reqs.length > 0) {
      report('Installing packages...', 85);
      await this.micropip.install(reqs);
    }

    // 5. Utility functions + Cognite client
    report('Setting up environment...', 90);
    this.pyodide.runPython(PYTHON_UTILS);

    report('Initializing Cognite client...', 95);
    this.pyodide.runPython(`
import os
os.environ["COGNITE_PROJECT"] = "${sdk.project}"
os.environ["COGNITE_BASE_URL"] = "${sdk.baseUrl}"
os.environ["COGNITE_TOKEN"] = "${sdk.token}"
os.environ["COGNITE_FUSION_NOTEBOOK"] = "1"
os.environ["MPLBACKEND"] = "AGG"
from cognite.client import CogniteClient
client = CogniteClient()
`);

    this._initialized = true;
    report('Ready', 100);
  }

  /** Execute Python code asynchronously. PyProxy results are converted to JSON-safe values. */
  async runCodeAsync(code: string): Promise<unknown> {
    const pyodide = this.requirePyodide();
    const raw = await pyodide.runPythonAsync(code);
    return this.toJsonSafe(raw);
  }

  /** Refresh the Cognite SDK token (e.g. after token rotation). */
  refreshToken(token: string): void {
    this.requirePyodide().runPython(
      `import os; os.environ["COGNITE_TOKEN"] = "${token}"`,
    );
  }

  private requirePyodide(): PyodideInstance {
    if (!this.pyodide) {
      throw new Error(
        'PyodideRuntime not initialized — call initialize() first',
      );
    }
    return this.pyodide;
  }

  /** Convert a Pyodide result to a JSON-safe JS value. */
  private toJsonSafe(value: unknown): unknown {
    if (value == null) return undefined;
    if (!isPyProxy(value)) return value;

    const pyodide = this.requirePyodide();
    const converter = pyodide.globals.get('as_json_string') as
      | ((obj: unknown) => string)
      | undefined;

    if (!converter) {
      throw new Error(
        'as_json_string not available — was initialize() called?',
      );
    }

    try {
      return JSON.parse((converter as (obj: unknown) => string)(value));
    } finally {
      destroyIfPyProxy(converter);
      destroyIfPyProxy(value);
    }
  }
}

// ============================================================================
// Singleton
// ============================================================================

let globalInstance: PyodideRuntime | undefined;

/**
 * Get or create the global PyodideRuntime singleton.
 * Config is only used on first call — subsequent calls return the existing instance.
 */
export function getGlobalPyodideRuntime(
  config: PyodideRuntimeConfig,
): PyodideRuntime {
  if (!globalInstance) {
    globalInstance = new PyodideRuntime(config);
  }
  return globalInstance;
}

/** Reset the global runtime (e.g. on logout). */
export function resetGlobalPyodideRuntime(): void {
  globalInstance = undefined;
}

/**
 * Pyodide entry point — browser-based Python execution for Atlas agents.
 *
 * Import path: @cognite/dune-utils/atlas-agent/pyodide
 */

// Runtime
export {
  PyodideRuntime,
  getGlobalPyodideRuntime,
  resetGlobalPyodideRuntime,
  clearPyodideCache,
} from './pyodide-runtime';
export type { PyodideRuntimeConfig, PyodideSDKConfig, PyodideInstance } from './pyodide-runtime';

// React hook
export { usePyodideRuntime } from './pyodide-react';
export type {
  PyodideProgress,
  UsePyodideRuntimeOptions,
  UsePyodideRuntimeReturn,
} from './pyodide-react';

/**
 * Python execution primitives for Pyodide-based tool execution.
 *
 * buildWrapper  — combines tool code + args into a single runnable Python snippet
 * formatOutput  — serialises the Pyodide result into a string for the agent
 */

/**
 * Unicode-safe base64 encoding.
 * Encodes to UTF-8 bytes first so non-ASCII characters survive round-tripping.
 */
function toBase64(str: string): string {
  const bytes = new TextEncoder().encode(str);
  let binary = '';
  for (const b of bytes) {
    binary += String.fromCharCode(b);
  }
  return btoa(binary);
}

/**
 * Build a Python wrapper that loads the tool code, parses base64-encoded args,
 * and calls handle(**args). Supports both sync and async handle functions.
 */
export function buildWrapper(code: string, argsJson: string): string {
  const encoded = toBase64(argsJson);
  return `
import json, base64, inspect
_args = json.loads(base64.b64decode("${encoded}").decode("utf-8"))
${code}
async def _exec():
    if "handle" not in globals():
        return {"_error": "No handle() function found in tool code"}
    return await handle(**_args) if inspect.iscoroutinefunction(handle) else handle(**_args)
_r = await _exec()
json.dumps(_r) if _r is not None and not isinstance(_r, str) else _r
`.trimStart();
}

/** Stringify Python result into tool output text. */
export function formatOutput(raw: unknown): string {
  if (raw == null) return '';
  if (typeof raw === 'string') return raw;
  return JSON.stringify(raw);
}

/**
 * useAtlasChat — plug-and-play React hook for Atlas agent conversations.
 *
 * Manages session lifecycle, message state, streaming, and abort support.
 * Separate entry point from core for tree-shaking.
 */

import { useState, useCallback, useRef, useEffect, useMemo } from 'react';
import type { CogniteClient } from '@cognite/sdk';
import { AtlasSession } from './session';
import type { AtlasTool, AtlasResponse, PythonRuntime, ToolCall } from './types';

// ============================================================================
// Types
// ============================================================================

export interface ChatMessage<TContext = unknown> {
  id: string;
  role: 'user' | 'assistant';
  text: string;
  timestamp: Date;
  isStreaming?: boolean;
  /** Tool calls (client-side and server-side) attached to this message */
  toolCalls?: ToolCall[];
  /** App-specific context data, populated via onResponse */
  context?: TContext;
}

export interface UseAtlasChatOptions<TContext = unknown> {
  client: CogniteClient | null;
  agentExternalId: string;
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  tools?: AtlasTool<any, any>[];
  /** Opt-in Python runtime (e.g. Pyodide) — only needed for agents that use Python tools. */
  pythonRuntime?: PythonRuntime;
  /** Messages to show on initial render (e.g. welcome message) */
  initialMessages?: ChatMessage<TContext>[];
  /** Called when a full response is received. Return context to merge into the assistant message. */
  onResponse?: (response: AtlasResponse) => TContext | void;
  /** Called before each send to inject app-level context (e.g. current todo state) into the request. */
  getAppContext?: () => string | undefined;
}

export interface UseAtlasChatReturn<TContext = unknown> {
  /** All messages in the conversation */
  messages: ChatMessage<TContext>[];
  /** Send a user message — automatically creates user + assistant messages, handles streaming */
  send: (text: string) => Promise<void>;
  /** True while the agent is responding */
  isStreaming: boolean;
  /** Current progress text (e.g. "Agent thinking", "Executing: render_widget") */
  progress: string | null;
  /** Error message if last send failed */
  error: string | null;
  /** Clear all messages and reset the session */
  reset: () => void;
  /** Cancel the current streaming response */
  abort: () => void;
  /** Replace messages (e.g. loading conversation history) */
  setMessages: (messages: ChatMessage<TContext>[]) => void;
}

// ============================================================================
// Helpers
// ============================================================================

let messageCounter = 0;

function generateId(): string {
  return `msg-${Date.now()}-${++messageCounter}`;
}

// ============================================================================
// Hook
// ============================================================================

export function useAtlasChat<TContext = unknown>(
  options: UseAtlasChatOptions<TContext>,
): UseAtlasChatReturn<TContext> {
  const { client, agentExternalId, tools, pythonRuntime, initialMessages, onResponse, getAppContext } = options;

  const [messages, setMessages] = useState<ChatMessage<TContext>[]>(initialMessages ?? []);
  const [isStreaming, setIsStreaming] = useState(false);
  const [progress, setProgress] = useState<string | null>(null);
  const [error, setError] = useState<string | null>(null);

  const sessionRef = useRef<AtlasSession | null>(null);
  const abortRef = useRef<AbortController | null>(null);
  const agentExternalIdRef = useRef(agentExternalId);
  const toolsRef = useRef(tools);
  const pythonRuntimeRef = useRef(pythonRuntime);
  const getAppContextRef = useRef(getAppContext);

  // Keep refs updated (array/object identity may change between renders)
  toolsRef.current = tools;
  pythonRuntimeRef.current = pythonRuntime;
  getAppContextRef.current = getAppContext;

  // Stable wrapper — always delegates to the latest getAppContext via ref.
  // Passed to AtlasSession once at creation so the session is never stale.
  const stableGetAppContext = useMemo(
    () => () => getAppContextRef.current?.(),
    [],
  );

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      abortRef.current?.abort();
    };
  }, []);

  const getSession = useCallback((): AtlasSession | null => {
    if (!client) return null;

    if (!sessionRef.current || agentExternalIdRef.current !== agentExternalId) {
      sessionRef.current = new AtlasSession({
        client,
        agentExternalId,
        tools: toolsRef.current,
        pythonRuntime: pythonRuntimeRef.current,
        getAppContext: stableGetAppContext,
      });
      agentExternalIdRef.current = agentExternalId;
    }

    return sessionRef.current;
  }, [client, agentExternalId]);

  const send = useCallback(
    async (text: string) => {
      const session = getSession();
      if (!session || isStreaming) return;

      setError(null);
      setIsStreaming(true);
      setProgress('Agent thinking');

      // Add user message
      const userMessage: ChatMessage<TContext> = {
        id: generateId(),
        role: 'user',
        text,
        timestamp: new Date(),
      };

      const assistantId = generateId();
      let accumulatedText = '';
      let assistantCreated = false;

      setMessages((prev) => [...prev, userMessage]);

      const abortController = new AbortController();
      abortRef.current = abortController;

      // ---- Helpers scoped to this send() call ----

      /** Update a single message by id */
      const updateMsg = (id: string, updates: Partial<ChatMessage<TContext>>) => {
        setMessages((prev) =>
          prev.map((m) => (m.id === id ? { ...m, ...updates } : m)),
        );
      };

      /** Finalize the assistant message — update if already created, otherwise add a new one */
      const finalizeAssistant = (fields: Partial<ChatMessage<TContext>>) => {
        if (assistantCreated) {
          updateMsg(assistantId, { isStreaming: false, ...fields });
        } else {
          setMessages((prev) => [
            ...prev,
            {
              id: assistantId,
              role: 'assistant' as const,
              timestamp: new Date(),
              text: '',
              isStreaming: false,
              ...fields,
            },
          ]);
        }
      };

      try {
        const response = await session.send(
          text,
          {
            onProgress: (progressText) => {
              setProgress(progressText);
            },
            onChunk: (chunk) => {
              if (!assistantCreated) {
                assistantCreated = true;
                setMessages((prev) => [
                  ...prev,
                  {
                    id: assistantId,
                    role: 'assistant' as const,
                    text: chunk,
                    timestamp: new Date(),
                    isStreaming: true,
                  },
                ]);
              }
              accumulatedText += chunk;
              updateMsg(assistantId, { text: accumulatedText });
            },
            onToolStart: (toolName) => {
              setProgress(`Executing: ${toolName}`);
            },
          },
          abortController.signal,
        );

        // Finalize assistant message
        finalizeAssistant({
          text:
            response.text ||
            (assistantCreated
              ? undefined
              : "I apologize, but I couldn't generate a response. Please try again."),
          toolCalls:
            response.toolCalls.length > 0
              ? response.toolCalls
              : undefined,
        });

        // Let the app attach context (e.g. applications) to the message
        const ctx = onResponse?.(response);
        if (ctx !== undefined) {
          updateMsg(assistantId, { context: ctx });
        }
      } catch (err) {
        if ((err as Error).name === 'AbortError') {
          // Cancelled by user — finalize any in-progress message
          if (assistantCreated) {
            updateMsg(assistantId, { isStreaming: false });
          }
        } else {
          const errorText =
            err instanceof Error ? err.message : 'Unknown error';
          setError(errorText);
          finalizeAssistant({ text: `Error: ${errorText}` });
        }
      } finally {
        setIsStreaming(false);
        setProgress(null);
        abortRef.current = null;
      }
    },
    [getSession, isStreaming, onResponse],
  );

  const reset = useCallback(() => {
    abortRef.current?.abort();
    setMessages(initialMessages ?? []);
    setIsStreaming(false);
    setProgress(null);
    setError(null);
    sessionRef.current = null;
  }, [initialMessages]);

  const abort = useCallback(() => {
    abortRef.current?.abort();
  }, []);

  return {
    messages,
    send,
    isStreaming,
    progress,
    error,
    reset,
    abort,
    setMessages,
  };
}

import { describe, it, expect, vi } from 'vitest';
import { AtlasSession } from './session';
import type { AtlasSessionConfig, RawAgentResponse } from './types';

/**
 * Minimal mock that satisfies CogniteClient just enough for AtlasSession.
 * AtlasClient.post is the only call path we exercise, so we stub it via the
 * prototype after construction.
 */
function createMockConfig(
  overrides?: Partial<AtlasSessionConfig>,
): AtlasSessionConfig {
  return {
    client: {} as AtlasSessionConfig['client'],
    agentExternalId: 'test-agent',
    ...overrides,
  };
}

/** Build a raw response with a tool action that requests a client tool call. */
function responseWithToolAction(
  actionId: string,
  toolName: string,
  args: Record<string, unknown>,
): RawAgentResponse {
  return {
    agentId: 'test-agent',
    agentExternalId: 'test-agent',
    response: {
      type: 'result',
      cursor: 'cursor-1',
      messages: [
        {
          role: 'assistant',
          actions: [
            {
              type: 'clientTool',
              actionId,
              clientTool: { name: toolName, arguments: args },
            },
          ],
        },
      ],
    },
  };
}

/** Build a terminal response (no actions). */
function terminalResponse(text: string): RawAgentResponse {
  return {
    agentId: 'test-agent',
    agentExternalId: 'test-agent',
    response: {
      type: 'result',
      cursor: 'cursor-2',
      messages: [{ role: 'assistant', content: { type: 'text', text } }],
    },
  };
}

describe(AtlasSession.name, () => {
  let postSpy: ReturnType<typeof vi.fn>;

  function createSession(config?: Partial<AtlasSessionConfig>): AtlasSession {
    const session = new AtlasSession(createMockConfig(config));

    // Stub the internal client.post so we never hit the network.
    postSpy = vi.fn();
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    (session as any).client = { post: postSpy };

    return session;
  }

  describe('appContext in continuation turns', () => {
    it('includes contextInformation on the initial user message', async () => {
      const session = createSession({
        getAppContext: () => 'todo state here',
      });

      postSpy.mockResolvedValueOnce(terminalResponse('Done'));

      await session.send('hello');

      const payload = postSpy.mock.calls[0][0];
      expect(payload.contextInformation).toEqual({
        appContext: 'todo state here',
      });
    });

    it('includes contextInformation on continuation turns after tool execution', async () => {
      let callCount = 0;
      const session = createSession({
        getAppContext: () => {
          callCount++;
          return `context-v${callCount}`;
        },
        tools: [
          {
            name: 'TestTool',
            description: 'test',
            parameters: { type: 'object', properties: {} },
            execute: () => ({ output: 'ok' }),
          },
        ],
      });

      // Turn 1: agent requests a tool call
      postSpy.mockResolvedValueOnce(
        responseWithToolAction('action-1', 'TestTool', {}),
      );
      // Turn 2: terminal response
      postSpy.mockResolvedValueOnce(terminalResponse('All done'));

      await session.send('do something');

      // First call: initial payload
      expect(postSpy).toHaveBeenCalledTimes(2);
      const initialPayload = postSpy.mock.calls[0][0];
      expect(initialPayload.contextInformation).toEqual({
        appContext: 'context-v1',
      });

      // Second call: continuation payload after tool execution
      const continuationPayload = postSpy.mock.calls[1][0];
      expect(continuationPayload.contextInformation).toEqual({
        appContext: 'context-v2',
      });
    });

    it('omits contextInformation when getAppContext returns undefined', async () => {
      const session = createSession({
        getAppContext: () => undefined,
      });

      postSpy.mockResolvedValueOnce(terminalResponse('Done'));

      await session.send('hello');

      const payload = postSpy.mock.calls[0][0];
      expect(payload.contextInformation).toBeUndefined();
    });

    it('omits contextInformation when getAppContext is not provided', async () => {
      const session = createSession();

      postSpy.mockResolvedValueOnce(terminalResponse('Done'));

      await session.send('hello');

      const payload = postSpy.mock.calls[0][0];
      expect(payload.contextInformation).toBeUndefined();
    });
  });
});

/**
 * AtlasSession — stateful conversation with validated tool execution.
 *
 * Manages the cursor, converts AtlasTool[] to API actions format,
 * validates tool arguments with TypeBox Value, and runs the tool execution loop.
 */

import { AtlasClient } from './client';
import { buildWrapper, formatOutput } from './python';
import { validateToolArguments } from './validation';
import type {
  AtlasTool,
  AtlasToolResult,
  AtlasResponse,
  AtlasSessionConfig,
  AgentToolConfig,
  PythonRuntime,
  StreamCallbacks,
  ChatPayload,
  ApiToolDefinition,
  RawAction,
  RawClientToolAction,
  RawAgentResponse,
  RequestMessage,
  ClientToolActionMessage,
  ToolCall,
} from './types';

const MAX_TURNS = 50;

/**
 * Parse raw arguments from the API (string or object) into a plain object.
 */
function parseArguments(
  raw: string | Record<string, unknown>,
): Record<string, unknown> {
  if (typeof raw === 'string') {
    try {
      return JSON.parse(raw);
    } catch {
      return {};
    }
  }
  return raw || {};
}

/**
 * Extract all actions from all messages in a raw response.
 */
function extractActions(raw: RawAgentResponse): RawAction[] {
  return raw.response.messages.flatMap((msg) => msg.actions ?? []);
}

/**
 * Extract server-side tool calls from reasoning blocks in a raw response.
 * Shape: messages[n].reasoning[n].data[n].toolCall
 */
function extractServerToolCalls(raw: RawAgentResponse): ToolCall[] {
  const calls: ToolCall[] = [];
  for (const msg of raw.response.messages) {
    for (const entry of msg.reasoning ?? []) {
      const data = (entry as { data?: unknown[] }).data;
      if (!Array.isArray(data)) continue;
      for (const item of data) {
        const tc = (item as { toolCall?: Record<string, unknown> }).toolCall;
        if (!tc) continue;
        const result = tc.result as Record<string, unknown> | undefined;
        calls.push({
          name: String(tc.name ?? ''),
          toolType: tc.toolType !== undefined ? String(tc.toolType) : undefined,
          input: tc.input,
          output: result?.output !== undefined ? String(result.output) : undefined,
          details: result?.result,
        });
      }
    }
  }
  return calls;
}

/**
 * Execute a server-declared Python tool via the pythonRuntime.
 */
async function executePythonTool(
  action: RawClientToolAction,
  toolConfig: AgentToolConfig,
  pythonRuntime: PythonRuntime | undefined,
  callbacks?: StreamCallbacks,
): Promise<{ result: AtlasToolResult; followup: ClientToolActionMessage }> {
  const toolName = action.clientTool.name;
  const pythonCode = String(toolConfig.configuration?.pythonCode ?? '').trim();

  const fail = (msg: string) => {
    const result: AtlasToolResult = { output: msg };
    callbacks?.onToolEnd?.(toolName, result);
    return { result, followup: createActionReply(action.actionId, result.output) };
  };

  if (!pythonCode) {
    return fail(`ERROR: pythonCode is empty in tool configuration for "${toolName}"`);
  }
  if (!pythonRuntime) {
    return fail(`ERROR: pythonRuntime is required to execute Python tool "${toolName}" but was not provided`);
  }

  try {
    const argsJson = JSON.stringify(parseArguments(action.clientTool.arguments));
    const wrapper = buildWrapper(pythonCode, argsJson);
    const raw = await pythonRuntime.runCodeAsync(wrapper);
    const result: AtlasToolResult = { output: formatOutput(raw) };
    callbacks?.onToolEnd?.(toolName, result);
    return { result, followup: createActionReply(action.actionId, result.output) };
  } catch (err) {
    return fail(`ERROR: ${err instanceof Error ? err.message : String(err)}`);
  }
}

/**
 * Execute a client tool action: validate args, run the tool, return result + followup.
 *
 * Dispatch order:
 *   1. Pre-registered JS tool in the tools Map (client-declared tools)
 *   2. Server-defined Python tool — fetch config from agent API, run via pythonRuntime
 */
async function executeClientTool(
  action: RawClientToolAction,
  tools: Map<string, AtlasTool>,
  fetchToolConfig: (name: string) => Promise<AgentToolConfig | null>,
  pythonRuntime: PythonRuntime | undefined,
  callbacks?: StreamCallbacks,
): Promise<{ result: AtlasToolResult; followup: ClientToolActionMessage }> {
  const toolName = action.clientTool.name;

  callbacks?.onToolStart?.(toolName);

  // 1. Pre-registered JS tool
  const tool = tools.get(toolName);
  if (tool) {
    const args = parseArguments(action.clientTool.arguments);
    let validatedArgs: unknown;
    try {
      validatedArgs = validateToolArguments(toolName, tool.parameters, args);
    } catch (err) {
      const errorOutput = err instanceof Error ? err.message : String(err);
      const result: AtlasToolResult = { output: `ERROR: ${errorOutput}` };
      callbacks?.onToolEnd?.(toolName, result);
      return { result, followup: createActionReply(action.actionId, result.output) };
    }
    const result = await tool.execute(validatedArgs);
    callbacks?.onToolEnd?.(toolName, result);
    return { result, followup: createActionReply(action.actionId, result.output) };
  }

  // 2. Server-declared tool — look up config from agent API
  const toolConfig = await fetchToolConfig(toolName);
  if (toolConfig?.type === 'runPythonCode') {
    return executePythonTool(action, toolConfig, pythonRuntime, callbacks);
  }

  const result: AtlasToolResult = { output: `Unknown client tool: ${toolName}` };
  callbacks?.onToolEnd?.(toolName, result);
  return { result, followup: createActionReply(action.actionId, result.output) };
}

/**
 * Stateful conversation session with validated tool execution.
 */
export class AtlasSession {
  private cursor?: string;
  private readonly client: AtlasClient;
  private readonly agentExternalId: string;
  private readonly tools: Map<string, AtlasTool>;
  private readonly apiActionsOrUndefined: ApiToolDefinition[] | undefined;
  private readonly pythonRuntime: PythonRuntime | undefined;
  private readonly getAppContext: (() => string | undefined) | undefined;
  /** Cached tool configs fetched from the agent API (populated lazily on first Python tool call). */
  private cachedAgentTools: AgentToolConfig[] | undefined;

  constructor(config: AtlasSessionConfig) {
    this.client = new AtlasClient(config.client);
    this.agentExternalId = config.agentExternalId;
    this.tools = new Map((config.tools || []).map((t) => [t.name, t]));
    this.pythonRuntime = config.pythonRuntime;
    this.getAppContext = config.getAppContext;

    // Inline toApiToolDefinition (only used here)
    const apiActions: ApiToolDefinition[] = (config.tools || []).map((tool) => ({
      type: 'clientTool' as const,
      clientTool: {
        name: tool.name,
        description: tool.description,
        parameters: tool.parameters,
      },
    }));
    this.apiActionsOrUndefined = apiActions.length > 0 ? apiActions : undefined;
  }

  /**
   * Send a user message. Handles the full tool execution loop internally.
   */
  async send(
    message: string,
    callbacks?: StreamCallbacks,
    signal?: AbortSignal,
  ): Promise<AtlasResponse> {
    const allToolCalls: ToolCall[] = [];

    const appContext = this.getAppContext?.();
    let payload: ChatPayload = {
      agentExternalId: this.agentExternalId,
      messages: [{ role: 'user', content: { type: 'text', text: message } }],
      actions: this.apiActionsOrUndefined,
      stream: true,
      ...(this.cursor && { cursor: this.cursor }),
      ...(appContext && { contextInformation: { appContext } }),
    };

    for (let turn = 0; turn < MAX_TURNS; turn++) {
      const raw = await this.client.post(payload, this.agentExternalId, callbacks, signal);
      const response = raw.response;

      if (response.type !== 'result') {
        throw new Error(`Unexpected response type: ${response.type}`);
      }

      if (response.cursor) {
        this.cursor = response.cursor;
      }

      // Collect server-side tool calls from reasoning blocks in this turn
      allToolCalls.push(...extractServerToolCalls(raw));

      const actions = extractActions(raw);

      // No actions → conversation turn is done
      if (actions.length === 0) {
        const text = response.messages?.[0]?.content?.text || '';
        return { text, cursor: this.cursor, toolCalls: allToolCalls, raw };
      }

      // Execute actions and build follow-up messages
      const followups: RequestMessage[] = [];

      for (const action of actions) {
        if (action.type === 'clientTool') {
          const { result, followup } = await executeClientTool(
            action,
            this.tools,
            (name) => this.fetchToolConfig(name),
            this.pythonRuntime,
            callbacks,
          );
          allToolCalls.push({ name: action.clientTool.name, output: result.output, details: result.details });
          followups.push(followup);
        } else if (action.type === 'toolConfirmation') {
          const toolName = action.toolConfirmation?.toolName;
          if (toolName) callbacks?.onToolStart?.(toolName);
          followups.push({
            role: 'action',
            type: 'toolConfirmation',
            actionId: action.actionId,
            status: 'ALLOW',
          });
        }
      }

      if (followups.length === 0) {
        const text = response.messages?.[0]?.content?.text || '';
        return { text, cursor: this.cursor, toolCalls: allToolCalls, raw };
      }

      // Prepare the next turn — re-evaluate appContext so it reflects state changes from tool execution
      const updatedAppContext = this.getAppContext?.();
      payload = {
        agentExternalId: this.agentExternalId,
        messages: followups,
        actions: this.apiActionsOrUndefined,
        stream: true,
        cursor: this.cursor,
        ...(updatedAppContext && { contextInformation: { appContext: updatedAppContext } }),
      };
    }

    throw new Error(`Max tool execution turns reached (${MAX_TURNS})`);
  }

  /**
   * Fetch a tool's config from the agent API (lazy, cached per session).
   * Used as a fallback when a clientTool action arrives for an unregistered tool.
   */
  private async fetchToolConfig(toolName: string): Promise<AgentToolConfig | null> {
    if (!this.cachedAgentTools) {
      const agent = await this.client.getAgentById(this.agentExternalId);
      this.cachedAgentTools = agent?.tools ?? [];
    }
    return this.cachedAgentTools.find((t) => t.name === toolName) ?? null;
  }

  /** Reset the session cursor (start a fresh conversation). */
  reset(): void {
    this.cursor = undefined;
  }

  /** Get the current cursor value. */
  getCursor(): string | undefined {
    return this.cursor;
  }

  /** Set the cursor (e.g. when restoring a conversation). */
  setCursor(cursor: string): void {
    this.cursor = cursor;
  }
}

function createActionReply(
  actionId: string,
  text: string,
): ClientToolActionMessage {
  return {
    role: 'action',
    type: 'clientTool',
    actionId,
    content: { type: 'text', text },
    data: [],
  };
}

/**
 * Core types for the Atlas Agent client library.
 *
 * This module is self-contained — no imports from outside the library
 * except external packages (@sinclair/typebox, @cognite/sdk).
 */

import type { TSchema, Static } from '@sinclair/typebox';
import type { CogniteClient } from '@cognite/sdk';

// ============================================================================
// Agent Types
// ============================================================================

/** Configuration for a tool stored in the agent's CDF config. */
export interface AgentToolConfig {
  name: string;
  type: string;
  configuration?: {
    pythonCode?: string;
    [key: string]: unknown;
  };
  [key: string]: unknown;
}

export interface Agent {
  externalId: string;
  name: string;
  description?: string;
  model?: string;
  instructions?: string;
  ownerId?: string;
  tools?: AgentToolConfig[];
  createdTime?: number;
  lastUpdatedTime?: number;
}

// ============================================================================
// Tool Types
// ============================================================================

/** Result from executing a tool */
export interface AtlasToolResult<TDetails = unknown> {
  /** Text sent back to the agent as tool output */
  output: string;
  /** Structured data for the app/UI */
  details?: TDetails;
}

/**
 * A client-side tool the Atlas agent can invoke.
 * TypeBox schema for type-safe params + runtime validation via TypeBox Value.
 */
export interface AtlasTool<
  TParameters extends TSchema = TSchema,
  TDetails = unknown,
> {
  name: string;
  description: string;
  parameters: TParameters;
  execute: (
    args: Static<TParameters>,
  ) => AtlasToolResult<TDetails> | Promise<AtlasToolResult<TDetails>>;
}

/** Minimal interface for executing Python code (e.g. Pyodide). */
export interface PythonRuntime {
  runCodeAsync(code: string): Promise<unknown>;
}

// ============================================================================
// Callback Types
// ============================================================================

export interface StreamCallbacks {
  onProgress?: (text: string) => void;
  onChunk?: (text: string) => void;
  onToolStart?: (toolName: string) => void;
  onToolEnd?: (toolName: string, result: AtlasToolResult) => void;
}

// ============================================================================
// Response Types (app-facing)
// ============================================================================

/** A single tool invocation (client-side or server-side), ready for the UI. */
export interface ToolCall {
  /** Friendly display name, e.g. "Find files" */
  name: string;
  /** Server-side tool type, e.g. "queryKnowledgeGraph" */
  toolType?: string;
  /** Raw input arguments */
  input?: unknown;
  /** Text returned to the agent as tool output */
  output?: string;
  /** Structured data for UI rendering */
  details?: unknown;
}

export interface AtlasResponse {
  text: string;
  cursor?: string;
  toolCalls: ToolCall[];
  raw: RawAgentResponse;
}

// ============================================================================
// Config Types
// ============================================================================

export interface AtlasSessionConfig {
  client: CogniteClient;
  agentExternalId: string;
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  tools?: AtlasTool<any, any>[];
  /** Opt-in Python runtime (e.g. Pyodide) — required only for agents that use Python tools. */
  pythonRuntime?: PythonRuntime;
  /** Called before each send to inject app-level context (e.g. current todo state) into the request. */
  getAppContext?: () => string | undefined;
}

// ============================================================================
// API Shared Primitives (maps to cog_ai…session.common)
// ============================================================================

/** Maps to AgentContentDTO */
export interface AgentContent {
  type: string;
  text?: string;
}

/** Maps to InstanceIdDTO */
export interface InstanceId {
  space: string;
  externalId: string;
}

/** Maps to ViewDTO */
export interface View {
  space: string;
  externalId: string;
  version: string;
}

/** Maps to PropertyVal type alias */
export type PropertyVal =
  | string
  | number
  | boolean
  | null
  | Record<string, unknown>
  | (string | null)[]
  | (number | null)[]
  | (boolean | null)[]
  | (Record<string, unknown> | null)[];

// ============================================================================
// Agent Data Types (maps to AgentDataDTO)
// ============================================================================

/** Maps to InstanceDataDTO — the only variant we narrow on in app code. */
export interface InstanceData {
  type: 'instance';
  view: View;
  instanceId: InstanceId;
  properties?: Record<string, PropertyVal>;
}

/**
 * Data items attached to response messages.
 * Only `InstanceData` is typed — other variants pass through as plain objects.
 */
export type AgentData = InstanceData | (Record<string, unknown> & { type: string });

// ============================================================================
// Tool Definition Types (maps to CustomClientActionDTO)
// ============================================================================

/** Maps to clientToolParameters */
export interface ClientToolParameters {
  type: 'object';
  description?: string;
  properties?: Record<string, Record<string, unknown>>;
  required?: string[];
  propertyOrdering?: string[];
}

/** Maps to CustomClientActionDTO */
export interface ApiToolDefinition {
  type: 'clientTool';
  clientTool: {
    name: string;
    description?: string;
    parameters: ClientToolParameters;
  };
}

// ============================================================================
// Request Message Types (maps to RequestMessageDTO)
// ============================================================================

/** Maps to AgentChatMessageUserRequestDTO */
export interface UserRequestMessage {
  role: 'user';
  content: AgentContent;
}

/** Maps to ClientToolCallActionMessageDTO */
export interface ClientToolActionMessage {
  role: 'action';
  type: 'clientTool';
  actionId: string;
  content: AgentContent;
  data: Array<Record<string, unknown>>;
}

/** Maps to UserConfirmationResponseDTO */
export interface UserConfirmationMessage {
  role: 'action';
  type: 'toolConfirmation';
  actionId: string;
  status: 'ALLOW' | 'DENY';
}

/** Maps to UserSessionResponseDTO */
export interface UserSessionMessage {
  role: 'action';
  type: 'userSession';
  actionId: string;
  nonce: string;
}

/** Maps to RequestMessageDTO (discriminated union) */
export type RequestMessage =
  | UserRequestMessage
  | ClientToolActionMessage
  | UserConfirmationMessage
  | UserSessionMessage;

// ============================================================================
// Session Context (maps to AgentSessionContextDTO)
// ============================================================================

export interface AgentSessionContext {
  instanceSpaces?: string[];
  dataModels?: Array<Record<string, unknown>>;
  timeZone?: string;
  appContext?: string;
}

// ============================================================================
// Chat Payload (maps to AgentSessionRequest)
// ============================================================================

export interface ChatPayload {
  agentExternalId?: string;
  messages: RequestMessage[];
  actions?: ApiToolDefinition[];
  contextInformation?: AgentSessionContext;
  cursor?: string;
  stream: boolean;
}

// ============================================================================
// Raw Response Types
// ============================================================================

/** Response action: agent requests client to execute a tool */
export interface RawClientToolAction {
  type: 'clientTool';
  actionId: string;
  clientTool: {
    name: string;
    arguments: string | Record<string, unknown>;
  };
}

/** Response action: agent requests user confirmation */
export interface RawToolConfirmationAction {
  type: 'toolConfirmation';
  actionId: string;
  toolConfirmation?: {
    toolName?: string;
    toolType?: string;
    toolArguments?: Record<string, unknown>;
    toolDescription?: string;
    content?: AgentContent;
  };
}

/** Response action: agent requests user session */
export interface RawUserSessionAction {
  type: 'userSession';
  actionId: string;
}

/** Action from agent response (discriminated by `type`). Unknown types are skipped by the session loop. */
export type RawAction =
  | RawClientToolAction
  | RawToolConfirmationAction
  | RawUserSessionAction;

/** A message in an agent response */
export interface RawMessage {
  content?: AgentContent;
  role: string;
  data?: AgentData[];
  reasoning?: Array<Record<string, unknown>>;
  actions?: RawAction[];
}

export interface RawAgentResponse {
  agentId: string;
  agentExternalId: string;
  response: {
    type: string;
    cursor?: string;
    messages: RawMessage[];
  };
}

import { describe, it, expect } from 'vitest';
import { Type } from '@sinclair/typebox';
import { validateToolArguments } from './validation';

describe('validateToolArguments', () => {
  const schema = Type.Object({
    name: Type.String(),
    count: Type.Number(),
  });

  it('returns args unchanged when they already match the schema', () => {
    const args = { name: 'hello', count: 3 };
    expect(validateToolArguments('myTool', schema, args)).toEqual(args);
  });

  it('coerces string numbers to numeric type via Value.Convert', () => {
    const args = { name: 'hello', count: '42' };
    const result = validateToolArguments('myTool', schema, args) as { name: string; count: number };
    expect(result.count).toBe(42);
    expect(typeof result.count).toBe('number');
  });

  it('throws a formatted error when args cannot be coerced into the schema', () => {
    const args = { name: 'hello', count: 'not-a-number' };
    expect(() => validateToolArguments('myTool', schema, args)).toThrow(
      /Tool "myTool" received invalid arguments:/,
    );
  });

  it('throws mentioning the offending path when a required field is missing', () => {
    const args = { name: 'hello' };
    expect(() => validateToolArguments('myTool', schema, args)).toThrow(
      /Tool "myTool" received invalid arguments:/,
    );
  });
});

/**
 * Runtime validation for tool arguments using TypeBox Value (CSP-safe; no eval).
 */

import type { TSchema } from "@sinclair/typebox";
import { Value } from "@sinclair/typebox/value";

/**
 * Validate and coerce tool arguments against a TypeBox / JSON Schema.
 * Returns coerced args on success. Throws a formatted error on validation failure.
 */
export function validateToolArguments(toolName: string, schema: TSchema, args: unknown): unknown {
  const coerced = Value.Convert(schema, args);
  if (Value.Check(schema, coerced)) return coerced;

  const errors = [...Value.Errors(schema, coerced)]
    .map((e) => `${e.path || "/"} ${e.message}`)
    .join("; ");
  throw new Error(`Tool "${toolName}" received invalid arguments: ${errors}`);
}