makelore/electron/coding-runtime/pi/model-tools/model-tool-registry.ts

import type { CodingPluginToolDefinition } from '../../../../shared/coding-plugins';
import type { ModelToolDetailsV1, ModelWebSearchFailureV1 } from '../../../../shared/model-tools';
import type { ProviderAccount } from '../../../shared/providers/types';
import type {
  PiProviderDescriptor,
  PiProviderSelection,
  PiWorkerCredentialProjection,
} from '../provider-config';
import {
  ModelWebSearchError,
  createModelWebSearchAdapter,
  type FrozenSelectedModel,
  type ModelWebSearchAdapter,
} from './web-search';

const WEB_SEARCH_TOOL: CodingPluginToolDefinition = Object.freeze({
  name: 'web_search',
  label: 'Web search',
  description: [
    'Search the public internet with the currently selected model for current information, facts, and sources.',
    'Use agent_browser only for interactive browser debugging, authenticated pages, and UI actions.',
    'When this tool fails, report the failure. Do not use agent_browser as a fallback.',
  ].join(' '),
  capabilityId: 'model.web-search',
  operation: 'search',
  roles: ['parent'],
  mutation: 'read',
  projectWriteLease: false,
  permissions: [],
  inputSchema: Object.freeze({
    type: 'object',
    additionalProperties: false,
    required: ['query'],
    properties: {
      query: { type: 'string', minLength: 1, maxLength: 2_000 },
    },
  }),
});

export interface ModelToolInvocationContext {
  conversationId: string;
  workerGeneration: number;
  runId: string;
  resourceId: string;
}

export interface ModelToolInvocationResult {
  content: Array<{ type: 'text'; text: string }>;
  details: ModelToolDetailsV1;
}

export interface RegisterModelToolWorkerInput {
  conversationId: string;
  generation: number;
  account: ProviderAccount;
  descriptor: PiProviderDescriptor;
  selection: PiProviderSelection;
  credential: PiWorkerCredentialProjection;
}

export interface ModelToolWorkerRegistration {
  tools: readonly CodingPluginToolDefinition[];
  dispose(): void;
}

export interface ModelToolRegistryOptions {
  adapter?: ModelWebSearchAdapter;
}

export interface ModelToolRegistryPort {
  registerWorker(input: RegisterModelToolWorkerInput): ModelToolWorkerRegistration;
  invoke(
    toolName: string,
    context: ModelToolInvocationContext,
    input: unknown,
    signal?: AbortSignal,
  ): Promise<ModelToolInvocationResult>;
}

interface WorkerBinding {
  selectedModel: FrozenSelectedModel;
}

function resolvedHeaderValue(
  value: string,
  credential: PiWorkerCredentialProjection,
): string | null {
  if (!value.startsWith('$')) return value;
  return credential.env[value.slice(1)]?.trim() || null;
}

function freezeSelectedModel(input: RegisterModelToolWorkerInput): FrozenSelectedModel | null {
  if (
    input.account.id !== input.selection.accountId
    || input.descriptor.accountId !== input.selection.accountId
    || input.descriptor.runtimeProviderId !== input.selection.runtimeProviderId
  ) {
    return null;
  }
  const capability = input.account.metadata
    ?.worksSquareModelCapabilities
    ?.[input.selection.modelId]
    ?.webSearch;
  if (!capability || !input.descriptor.baseUrl) return null;
  const expectedApi = capability.adapter === 'bailian-chat-completions'
    ? 'openai-completions'
    : 'openai-responses';
  if (input.descriptor.api !== expectedApi) return null;

  const headers: Record<string, string> = {};
  for (const [name, rawValue] of Object.entries(input.descriptor.headers)) {
    const value = resolvedHeaderValue(rawValue, input.credential);
    if (!value) return null;
    headers[name] = value;
  }
  const hasAuthorization = Object.keys(headers).some((name) => name.toLowerCase() === 'authorization');
  const apiKey = input.descriptor.apiKeyEnv
    ? input.credential.env[input.descriptor.apiKeyEnv]?.trim()
    : undefined;
  if (!hasAuthorization && apiKey) headers.Authorization = `Bearer ${apiKey}`;

  return Object.freeze({
    accountId: input.account.id,
    runtimeProviderId: input.selection.runtimeProviderId,
    modelId: input.selection.modelId,
    generation: input.generation,
    baseUrl: input.descriptor.baseUrl,
    headers: Object.freeze(headers),
    capability: Object.freeze({ ...capability }),
  });
}

function failure(
  modelId: string,
  error: ModelWebSearchError,
): ModelToolInvocationResult {
  const details: ModelWebSearchFailureV1 = {
    schema: 'makelore-model-tool.v1',
    tool: 'web_search',
    status: 'failed',
    modelId,
    error: {
      code: error.code,
      message: error.message,
      httpStatus: error.status,
      retryable: error.retryable,
    },
  };
  return {
    content: [{ type: 'text', text: `${error.code}: ${error.message}` }],
    details,
  };
}

function queryFromInput(value: unknown): string | null {
  if (!value || typeof value !== 'object' || Array.isArray(value)) return null;
  const record = value as Record<string, unknown>;
  if (Object.keys(record).some((key) => key !== 'query')) return null;
  const query = typeof record.query === 'string' ? record.query.trim() : '';
  return query && query.length <= 2_000 ? query : null;
}

function successText(answer: string, sources: readonly { title: string; url: string }[]): string {
  if (sources.length === 0) return answer;
  return `${answer}\n\nSources:\n${sources.map(({ title, url }) => `- ${title}: ${url}`).join('\n')}`;
}

export class ModelToolRegistry implements ModelToolRegistryPort {
  private readonly adapter: ModelWebSearchAdapter;
  private readonly bindings = new Map<string, WorkerBinding>();

  constructor(options: ModelToolRegistryOptions = {}) {
    this.adapter = options.adapter ?? createModelWebSearchAdapter();
  }

  registerWorker(input: RegisterModelToolWorkerInput): ModelToolWorkerRegistration {
    const selectedModel = freezeSelectedModel(input);
    if (!selectedModel) return { tools: [], dispose: () => undefined };
    const binding = { selectedModel };
    this.bindings.set(input.conversationId, binding);
    return {
      tools: [WEB_SEARCH_TOOL],
      dispose: () => {
        if (this.bindings.get(input.conversationId) === binding) {
          this.bindings.delete(input.conversationId);
        }
      },
    };
  }

  async invoke(
    toolName: string,
    context: ModelToolInvocationContext,
    input: unknown,
    signal = new AbortController().signal,
  ): Promise<ModelToolInvocationResult> {
    const binding = this.bindings.get(context.conversationId);
    const modelId = binding?.selectedModel.modelId ?? 'unknown';
    if (toolName !== 'web_search') {
      return failure(modelId, new ModelWebSearchError(
        'model_web_search_unsupported',
        400,
        false,
        'The requested model tool is unavailable',
      ));
    }
    if (!binding || binding.selectedModel.generation !== context.workerGeneration) {
      return failure(modelId, new ModelWebSearchError(
        'model_context_changed',
        409,
        false,
        'The selected model context changed before Web Search completed',
      ));
    }
    const query = queryFromInput(input);
    if (!query) {
      return failure(modelId, new ModelWebSearchError(
        'model_web_search_invalid_result',
        502,
        false,
        'Web Search query must contain between 1 and 2000 characters',
      ));
    }
    try {
      const details = await this.adapter.search({
        query,
        selectedModel: binding.selectedModel,
        parentTurnId: context.runId,
        toolCallId: context.resourceId,
      }, signal);
      if (this.bindings.get(context.conversationId) !== binding) {
        return failure(modelId, new ModelWebSearchError(
          'model_context_changed',
          409,
          false,
          'The selected model context changed while Web Search was running',
        ));
      }
      return {
        content: [{ type: 'text', text: successText(details.answer, details.sources) }],
        details,
      };
    } catch (error) {
      const normalized = error instanceof ModelWebSearchError
        ? error
        : new ModelWebSearchError(
            'model_web_search_unavailable',
            503,
            true,
            'The selected model Web Search transport is unavailable',
          );
      return failure(modelId, normalized);
    }
  }
}