CopilotKitProvider.tsx

import React, { useEffect, useMemo, useRef, useState } from "react";
import type { ReactNode } from "react";
import {
  CopilotKitContext,
  LicenseContext,
} from "@copilotkit/react-core/v2/context";
import type {
  CopilotKitContextValue,
  CopilotKitCoreReact as CopilotKitCoreReactInstance,
} from "@copilotkit/react-core/v2/context";
import { CopilotKitCoreReact } from "@copilotkit/react-core/v2/headless";
import type { CopilotKitCoreErrorCode } from "@copilotkit/core";
import type { DebugConfig, RuntimeLicenseStatus } from "@copilotkit/shared";
import { createLicenseContextValue } from "@copilotkit/shared";
import { RenderToolProvider } from "./hooks/RenderToolContext";

export interface CopilotKitNativeProviderProps {
  children: ReactNode;
  /** URL of the CopilotKit runtime endpoint */
  runtimeUrl: string;
  /** Custom headers sent with every request */
  headers?: Record<string, string> | (() => Record<string, string>);
  /**
   * Credentials mode for fetch requests (e.g., "include" for HTTP-only cookies in cross-origin requests).
   */
  credentials?: RequestCredentials;
  /** Whether the runtime uses a single-route endpoint */
  useSingleEndpoint?: boolean;
  /** Custom properties forwarded to agents */
  properties?: Record<string, unknown>;
  /**
   * Error handler called when CopilotKit encounters an error.
   * Fires for all error types (runtime connection failures, agent errors, tool errors).
   * If not provided, errors are logged to console.error.
   */
  onError?: (event: {
    error: Error;
    code: CopilotKitCoreErrorCode;
    context: Record<string, any>;
  }) => void | Promise<void>;
  /**
   * Enable debug logging for the client-side event pipeline.
   * When `true`, enables verbose logging from the core instance.
   */
  debug?: DebugConfig;
  /**
   * Default throttle interval (ms) for `onMessagesChanged` / `onStateChanged`
   * subscriptions. Individual subscriptions can override with their own `throttleMs`.
   */
  defaultThrottleMs?: number;
  // Cloud features (publicApiKey, licenseToken) — not yet supported on React Native
}

/**
 * CopilotKit provider for React Native.
 *
 * A lightweight alternative to the web CopilotKitProvider that avoids
 * web-only dependencies (DOM, CSS, Radix UI, Lit, etc).
 *
 * Polyfills are auto-imported when `@copilotkit/react-native` is loaded,
 * so a separate `import "@copilotkit/react-native/polyfills"` is no longer
 * required (though it remains available for advanced use).
 *
 * Usage:
 * ```tsx
 * import { CopilotKitProvider } from "@copilotkit/react-native";
 *
 * function App() {
 *   return (
 *     <CopilotKitProvider runtimeUrl="https://your-runtime/api/copilotkit">
 *       <ChatScreen />
 *     </CopilotKitProvider>
 *   );
 * }
 * ```
 */
export const CopilotKitProvider: React.FC<CopilotKitNativeProviderProps> = ({
  children,
  runtimeUrl,
  headers: headersProp,
  credentials,
  useSingleEndpoint,
  properties,
  onError,
  debug,
  defaultThrottleMs,
}) => {
  // Resolve headers from function or static object (matches web provider pattern)
  const resolvedHeaders =
    typeof headersProp === "function" ? headersProp() : headersProp;

  // Stabilize headers/properties references to avoid effect churn when callers
  // pass inline object literals (e.g. headers={{}} or the undefined default).
  // eslint-disable-next-line react-hooks/exhaustive-deps
  const stableHeaders = useMemo(
    () => resolvedHeaders ?? {},
    [JSON.stringify(resolvedHeaders)],
  );
  // eslint-disable-next-line react-hooks/exhaustive-deps
  const stableProperties = useMemo(
    () => properties ?? {},
    [JSON.stringify(properties)],
  );

  const copilotkitRef = useRef<CopilotKitCoreReactInstance | null>(null);

  if (copilotkitRef.current === null) {
    const instance: CopilotKitCoreReactInstance = new CopilotKitCoreReact({
      runtimeUrl,
      runtimeTransport:
        useSingleEndpoint === true
          ? "single"
          : useSingleEndpoint === false
            ? "rest"
            : "auto",
      headers: stableHeaders,
      credentials,
      properties: stableProperties,
      debug,
    });
    // Set initial defaultThrottleMs synchronously so child hooks see the
    // correct value on their first render (before useEffect fires).
    if (defaultThrottleMs !== undefined) {
      instance.setDefaultThrottleMs(defaultThrottleMs);
    }
    copilotkitRef.current = instance;
  }

  const copilotkit = copilotkitRef.current;

  // Sync props to core instance
  useEffect(() => {
    copilotkit.setRuntimeUrl(runtimeUrl);
    copilotkit.setRuntimeTransport(
      useSingleEndpoint === true
        ? "single"
        : useSingleEndpoint === false
          ? "rest"
          : "auto",
    );
    copilotkit.setHeaders(stableHeaders);
    copilotkit.setCredentials(credentials);
    copilotkit.setProperties(stableProperties);
    copilotkit.setDebug(debug);
  }, [
    runtimeUrl,
    useSingleEndpoint,
    stableHeaders,
    credentials,
    stableProperties,
    debug,
    copilotkit,
  ]);

  // Sync defaultThrottleMs to the core instance on prop changes.
  // Initial value is set synchronously during instance creation (inside the
  // ref guard above), so this only handles subsequent updates.
  useEffect(() => {
    copilotkit.setDefaultThrottleMs(defaultThrottleMs);
  }, [copilotkit, defaultThrottleMs]);

  // Track executing tool call IDs at the provider level.
  // Critical for HITL reconnection: onToolExecutionStart fires before child
  // components mount, so we must capture the state here.
  const [executingToolCallIds, setExecutingToolCallIds] = useState<
    ReadonlySet<string>
  >(() => new Set());

  const [runtimeLicenseStatus, setRuntimeLicenseStatus] = useState<
    RuntimeLicenseStatus | undefined
  >(undefined);

  // Use ref to avoid subscription churn when onError changes
  const onErrorRef = useRef(onError);
  useEffect(() => {
    onErrorRef.current = onError;
  }, [onError]);

  // Single subscription for tool execution tracking and error handling.
  // Tool call IDs are tracked at the provider level because onToolExecutionStart
  // fires before child components mount — critical for HITL reconnection.
  useEffect(() => {
    const subscription = copilotkit.subscribe({
      onToolExecutionStart: ({ toolCallId }) => {
        setExecutingToolCallIds((prev) => {
          if (prev.has(toolCallId)) return prev;
          const next = new Set(prev);
          next.add(toolCallId);
          return next;
        });
      },
      onToolExecutionEnd: ({ toolCallId }) => {
        setExecutingToolCallIds((prev) => {
          if (!prev.has(toolCallId)) return prev;
          const next = new Set(prev);
          next.delete(toolCallId);
          return next;
        });
      },
      onError: (event) => {
        if (onErrorRef.current) {
          onErrorRef.current(event);
        } else {
          console.error(
            `[CopilotKit] Error (${event.code}):`,
            event.error,
            event.context ?? {},
          );
        }
      },
      onRuntimeConnectionStatusChanged: () => {
        setRuntimeLicenseStatus(copilotkit.licenseStatus);
      },
    });
    return () => subscription.unsubscribe();
  }, [copilotkit]);

  const contextValue: CopilotKitContextValue = useMemo(
    () => ({
      copilotkit,
      executingToolCallIds,
    }),
    [copilotkit, executingToolCallIds],
  );

  // License context — driven by server-reported status via /info endpoint
  const licenseContextValue = useMemo(
    () => createLicenseContextValue(runtimeLicenseStatus),
    [runtimeLicenseStatus],
  );

  return (
    <CopilotKitContext.Provider value={contextValue}>
      <LicenseContext.Provider value={licenseContextValue}>
        <RenderToolProvider>{children}</RenderToolProvider>
      </LicenseContext.Provider>
    </CopilotKitContext.Provider>
  );
};