React AG-UI streaming preview

// eslint-disable-next-line @nx/enforce-module-boundaries -- The isolated consumer copies this internal example layout.
import { ExampleFooter, ExampleLayout } from '../../../../shared/react-example';
import { useEffect, useState } from 'react';
import { useAgent } from '@threadplane/react';
import { ChatInput, TextTranscript } from '@threadplane/react/chat';
import {
  createConnectedApplication,
  type streamingConnection,
} from './connection';
 
export function StreamingDemo({
  connection,
  onReady,
}: {
  readonly connection: ReturnType<typeof streamingConnection>;
  readonly onReady: () => void;
}) {
  const [application] = useState(() => createConnectedApplication(connection));
  const [draft, setDraft] = useState('');
  const snapshot = useAgent(application);
  useEffect(() => {
    onReady();
  }, [onReady]);
  useEffect(
    () => () => {
      void application.dispose();
    },
    [application]
  );
  useEffect(() => {
    setDraft('');
  }, [snapshot.viewGeneration]);
  const status = snapshot.busy
    ? snapshot.activity === 'replacing'
      ? 'Starting new conversation…'
      : 'Receiving response…'
    : !snapshot.canSubmit
    ? 'Start a new conversation to continue.'
    : snapshot.outcome === 'success'
    ? 'Response complete.'
    : 'Ready.';
  return (
    <ExampleLayout className="streaming-demo">
      <header>
        <div>
          <h1>AG-UI streaming</h1>
        </div>
        <button
          type="button"
          disabled={snapshot.busy}
          onClick={() => void application.newConversation()}
        >
          New conversation
        </button>
      </header>
      <p>
        Watch an ordinary text reply arrive from the connected AG-UI agent. This
        experimental preview keeps the current conversation until you start a
        new one. Tools, pauses and uncertain outcomes require a new
        conversation.
      </p>
      <section aria-label="Runtime connection">
        <h2>Runtime connection</h2>
        <p>{connection.label}</p>
      </section>
      <TextTranscript messages={snapshot.rows} label="Conversation" />
      <ExampleFooter>
        {snapshot.error && <p role="alert">{snapshot.error}</p>}
        <ChatInput
          label="Message"
          placeholder="Ask a fictional streaming question…"
          value={draft}
          onValueChange={setDraft}
          disabled={!snapshot.canSubmit}
          busy={snapshot.busy}
          onStop={snapshot.busy ? () => void application.stop() : undefined}
          onSubmit={(text) => {
            if (!application.getSnapshot().canSubmit) return false;
            void application.submit(text);
            if (!application.getSnapshot().busy) return false;
            return true;
          }}
        />
        <p role="status">{status}</p>
      </ExampleFooter>
    </ExampleLayout>
  );
}

import { createSession } from '@threadplane/ag-ui';
import { createStreamingApplication } from './application';
// eslint-disable-next-line @nx/enforce-module-boundaries -- The isolated build copies this neutral internal bridge; it is not a published package.
import type { RuntimeBridgeConfiguration } from '../../../../../libs/cockpit-runtime-bridge/src/index';
export function streamingConnection(
  configuration: RuntimeBridgeConfiguration
): {
  url: string;
  label: 'Shared runtime' | 'Developer runtime';
} {
  if (
    configuration.status === 'error' ||
    (configuration.status === 'configured' &&
      configuration.target.kind === 'langsmith')
  )
    throw new Error('Runtime configuration failed.');
  const target =
    configuration.status === 'configured' &&
    configuration.target.kind === 'ag-ui'
      ? configuration.target
      : null;
  return {
    url:
      target?.endpoint ??
      new URL('/ag-ui/streaming/agent', window.location.origin).href,
    label: target ? 'Developer runtime' : 'Shared runtime',
  };
}
export function createConnectedApplication(
  connection: ReturnType<typeof streamingConnection>
) {
  return createStreamingApplication({
    sessionFactory(threadId) {
      const session = createSession({ threadId, url: connection.url });
      try {
        const snapshot = session.getSnapshot();
        if (
          !snapshot ||
          !Array.isArray(snapshot.transcript) ||
          !Array.isArray(snapshot.subagents) ||
          ![
            'getSnapshot',
            'subscribe',
            'submit',
            'resume',
            'stop',
            'dispose',
          ].every(
            (method) =>
              typeof session[method as keyof typeof session] === 'function'
          )
        )
          throw new Error('Runtime session is unavailable.');
      } catch {
        if (typeof session?.dispose === 'function') {
          try {
            void Promise.resolve(session.dispose()).catch(() => undefined);
          } catch {
            /* Rejected candidates have no authority. */
          }
        }
        throw new Error('Runtime session is unavailable.');
      }
      return session;
    },
  });
}

import type { Session, TextTranscriptRow } from '@threadplane/ag-ui';
 
export type NativeSnapshot = ReturnType<Session['getSnapshot']>;
export interface OwnedTurn {
  readonly threadId: string;
  readonly runId: string;
  readonly human: TextTranscriptRow;
  readonly prefix: readonly TextTranscriptRow[];
}
 
function ordinaryRows(
  transcript: NativeSnapshot['transcript'],
  valid: () => boolean
): readonly TextTranscriptRow[] | undefined {
  const rows: TextTranscriptRow[] = [];
  const identities = new Set<string>();
  for (const message of transcript) {
    if (!valid()) return;
    const { id, role, content, subagentRunId } = message;
    if (
      typeof id !== 'string' ||
      !id ||
      identities.has(id) ||
      (role !== 'user' && role !== 'assistant') ||
      typeof content !== 'string' ||
      !content ||
      typeof subagentRunId === 'string' ||
      (role === 'assistant' && !!message.toolCalls?.length) ||
      !valid()
    )
      return;
    identities.add(id);
    rows.push(Object.freeze({ id, role, content }));
  }
  return valid() ? Object.freeze(rows) : undefined;
}
 
function same(left: TextTranscriptRow, right: TextTranscriptRow) {
  return (
    left.id === right.id &&
    left.role === right.role &&
    left.content === right.content
  );
}
 
function read(snapshot: NativeSnapshot, valid: () => boolean) {
  if (!valid()) return;
  const status = snapshot.status;
  if (!valid()) return;
  const run = snapshot.run;
  if (!valid()) return;
  const decision = snapshot.decision;
  if (!valid()) return;
  const subagents = snapshot.subagents;
  if (!valid()) return;
  const transcript = snapshot.transcript;
  if (!valid() || decision || run?.legacyInterrupt || subagents.length) return;
  const rows = ordinaryRows(transcript, valid);
  return rows && valid() ? { status, run, rows } : undefined;
}
 
export function captureTurn(
  snapshot: NativeSnapshot,
  threadId: string,
  previousRunId: string | undefined,
  prefix: readonly TextTranscriptRow[],
  valid: () => boolean = () => true
): OwnedTurn | undefined {
  const state = read(snapshot, valid);
  if (
    !state ||
    state.status !== 'running' ||
    !state.run?.id ||
    state.run.id === previousRunId ||
    state.run.outcome !== undefined ||
    state.rows.length !== prefix.length + 1 ||
    !prefix.every((message, index) => same(message, state.rows[index])) ||
    !valid()
  )
    return;
  const human = state.rows.at(-1)!;
  if (human.role !== 'user' || !threadId || !valid()) return;
  return Object.freeze({
    threadId,
    runId: state.run.id,
    human,
    prefix: Object.freeze([...prefix]),
  });
}
 
export function confirmTranscript(
  snapshot: NativeSnapshot,
  turn: OwnedTurn,
  valid: () => boolean = () => true
): readonly TextTranscriptRow[] | undefined {
  const state = read(snapshot, valid);
  if (
    !state ||
    state.status !== 'idle' ||
    state.run?.id !== turn.runId ||
    state.run.outcome !== 'success' ||
    !valid()
  )
    return;
  const terminal = state.run.terminal;
  if (
    terminal?.type !== 'RUN_FINISHED' ||
    terminal.threadId !== turn.threadId ||
    terminal.runId !== turn.runId ||
    (terminal.outcome !== undefined && terminal.outcome.type !== 'success') ||
    (terminal.outcome?.type === 'success' &&
      !!terminal.outcome.pendingToolCallIds?.length) ||
    state.rows.length !== turn.prefix.length + 2 ||
    !turn.prefix.every((message, index) => same(message, state.rows[index])) ||
    !same(state.rows.at(-2)!, turn.human) ||
    state.rows.at(-1)!.role !== 'assistant' ||
    !valid()
  )
    return;
  return state.rows;
}

import {
  projectTextTranscript,
  type Session,
  type TextTranscriptRow,
} from '@threadplane/ag-ui';
import type { CompleteOutcome } from '@threadplane/core';
import {
  captureTurn,
  confirmTranscript,
  type NativeSnapshot,
  type OwnedTurn,
} from './canonical-transcript';
 
export interface StreamingSnapshot {
  readonly threadId: string | null;
  readonly native: NativeSnapshot | undefined;
  readonly rows: readonly TextTranscriptRow[];
  readonly busy: boolean;
  readonly canSubmit: boolean;
  readonly activity: 'idle' | 'running' | 'replacing';
  readonly outcome: CompleteOutcome | null;
  readonly error: string | null;
  readonly viewGeneration: number;
}
 
/** Owns lifetime and admission; the native session owns all protocol execution. */
export function createStreamingApplication(options: {
  readonly sessionFactory: (threadId: string) => Session;
}) {
  type Primary = {
    session: Session;
    threadId: string;
    closed: boolean;
    release?: () => void;
    unsafe: boolean;
    confirmed: readonly TextTranscriptRow[];
    previousRunId?: string;
    turn?: OwnedTurn;
  };
  const listeners = new Set<() => void>(),
    cleanups = new Set<Promise<void>>();
  let primary: Primary | undefined;
  let operation: AbortController | undefined;
  let disposed = false,
    disposal: Promise<void> | undefined;
  let snapshot: StreamingSnapshot = Object.freeze({
    threadId: null,
    native: undefined,
    rows: Object.freeze([]),
    busy: false,
    canSubmit: true,
    activity: 'idle',
    outcome: null,
    error: null,
    viewGeneration: 0,
  });
  function publish(update: Partial<StreamingSnapshot>) {
    if (disposed) return;
    snapshot = Object.freeze({ ...snapshot, ...update });
    for (const notify of listeners) {
      try {
        notify();
      } catch {
        /* A view cannot interrupt native execution or cleanup. */
      }
    }
  }
  function current(selected: Primary) {
    return (
      !disposed &&
      primary === selected &&
      !selected.closed &&
      !operation?.signal.aborted
    );
  }
  function close(selected: Primary) {
    if (selected.closed) return Promise.resolve();
    selected.closed = true;
    try {
      selected.release?.();
    } catch {
      /* Revoked observation cannot prevent native cleanup. */
    }
    const cleanup = Promise.resolve()
      .then(() => selected.session.dispose())
      .then(
        () => undefined,
        () => undefined
      );
    cleanups.add(cleanup);
    void cleanup.then(() => cleanups.delete(cleanup));
    return cleanup;
  }
  function detach() {
    const previous = primary;
    primary = undefined;
    return previous ? close(previous) : Promise.resolve();
  }
  function read(selected: Primary, valid: () => boolean) {
    if (!valid()) return;
    const root = selected.session.getSnapshot();
    if (!valid()) return;
    const transcript = root.transcript;
    if (!valid()) return;
    const subagents = root.subagents;
    if (!valid()) return;
    const decision = root.decision;
    if (!valid()) return;
    const run = root.run;
    if (!valid()) return;
    const status = root.status;
    if (!valid()) return;
    const blocked =
      status === 'error' ||
      !!decision ||
      !!subagents.length ||
      !!run?.legacyInterrupt ||
      (run?.terminal?.type === 'RUN_FINISHED' &&
        run.terminal.outcome?.type === 'success' &&
        !!run.terminal.outcome.pendingToolCallIds?.length) ||
      transcript.some(
        (message) =>
          (message.role !== 'user' && message.role !== 'assistant') ||
          typeof message.subagentRunId === 'string' ||
          (message.role === 'assistant' && !!message.toolCalls?.length)
      );
    return valid() ? { root, blocked } : undefined;
  }
  function capture(selected: Primary, valid: () => boolean) {
    const state = read(selected, valid);
    if (!state || !valid()) return;
    selected.unsafe ||= state.blocked;
    if (!selected.turn && snapshot.activity === 'running') {
      const turn = captureTurn(
        state.root,
        selected.threadId,
        selected.previousRunId,
        selected.confirmed,
        valid
      );
      if (!valid()) return;
      if (turn) selected.turn = turn;
    }
    const confirmed =
      selected.turn && !selected.unsafe
        ? confirmTranscript(state.root, selected.turn, valid)
        : undefined;
    if (!valid()) return;
    const rows = projectTextTranscript(state.root.transcript, snapshot.rows);
    return valid() ? { native: state.root, rows, confirmed } : undefined;
  }
  function observe(selected: Primary) {
    const valid = () => current(selected);
    try {
      const result = capture(selected, valid);
      if (!result || !valid()) return;
      publish({
        native: result.native,
        rows: result.rows,
        ...(selected.unsafe ? { canSubmit: false } : {}),
      });
    } catch {
      if (valid()) {
        selected.unsafe = true;
        publish({
          canSubmit: false,
          error:
            'The streaming result could not be confirmed. Start a new conversation.',
        });
      }
    }
  }
  function install(valid: () => boolean) {
    const threadId = crypto.randomUUID();
    const session = options.sessionFactory(threadId);
    const selected: Primary = {
      session,
      threadId,
      closed: false,
      unsafe: false,
      confirmed: [],
    };
    if (!valid()) {
      void close(selected);
      return;
    }
    primary = selected;
    const release = session.subscribe(() => observe(selected));
    selected.release = release;
    if (selected.closed) release();
    if (!valid() || primary !== selected) {
      void close(selected);
      return;
    }
    return selected;
  }
  function admit(activity: StreamingSnapshot['activity'], execution: boolean) {
    const admitted = new AbortController();
    operation = admitted;
    const valid = () =>
      !disposed && operation === admitted && !admitted.signal.aborted;
    publish({
      busy: true,
      canSubmit: false,
      activity,
      ...(execution ? { outcome: null, error: null } : {}),
    });
    return { admitted, valid };
  }
  function finish(admitted: AbortController) {
    if (!disposed && operation === admitted) {
      operation = undefined;
      publish({ busy: false, activity: 'idle' });
    }
  }
  async function execute(text: string) {
    if (disposed || snapshot.busy || !snapshot.canSubmit || !text.trim())
      return;
    const before = primary;
    const initialValid = () =>
      !disposed &&
      !snapshot.busy &&
      snapshot.canSubmit &&
      primary === before &&
      !before?.closed;
    let state: ReturnType<typeof read>;
    try {
      state = before && read(before, initialValid);
    } catch {
      if (initialValid())
        publish({
          canSubmit: false,
          error:
            'The streaming session is unavailable. Start a new conversation.',
        });
      return;
    }
    if (!initialValid()) return;
    if (before && (before.unsafe || state?.blocked)) {
      publish({ canSubmit: false });
      return;
    }
    const { admitted, valid } = admit('running', true);
    try {
      if (!valid()) return;
      const selected = primary ?? install(valid);
      if (!selected || !valid()) return;
      publish({ threadId: selected.threadId });
      const ownedValid = () => valid() && current(selected);
      const root = read(selected, ownedValid);
      if (!root || !ownedValid()) return;
      if (root.blocked || selected.unsafe) {
        selected.unsafe = true;
        publish({ canSubmit: false });
        return;
      }
      selected.previousRunId = root.root.run?.id;
      selected.turn = undefined;
      if (!ownedValid()) return;
      const pending = selected.session.submit(text, {
        signal: admitted.signal,
      });
      void pending.catch(() => undefined);
      if (!ownedValid()) return;
      observe(selected);
      const outcome = await pending;
      if (!ownedValid()) return;
      const result = capture(selected, ownedValid);
      if (!result || !ownedValid()) return;
      const safe =
        outcome === 'success' && !!result.confirmed && !selected.unsafe;
      if (safe) selected.confirmed = result.confirmed!;
      publish({
        native: result.native,
        rows: result.rows,
        outcome,
        canSubmit: safe,
        error: safe
          ? null
          : 'The current streaming outcome could not be confirmed. Start a new conversation.',
      });
    } catch {
      if (valid())
        publish({
          canSubmit: false,
          outcome: 'error',
          error: 'The AG-UI request failed. Start a new conversation.',
        });
    } finally {
      finish(admitted);
    }
  }
  async function newConversation() {
    if (disposed || snapshot.busy) return;
    const { admitted, valid } = admit('replacing', false);
    const cleanup = detach();
    try {
      if (!valid()) return;
      publish({
        threadId: null,
        native: undefined,
        rows: Object.freeze([]),
        outcome: null,
        error: null,
        viewGeneration: snapshot.viewGeneration + 1,
      });
      if (!valid()) return;
      await cleanup;
      if (valid()) publish({ canSubmit: true });
    } finally {
      finish(admitted);
    }
  }
  async function stop() {
    if (disposed || !operation) return;
    operation.abort();
    operation = undefined;
    const selected = primary,
      cleanup = detach();
    publish({
      busy: false,
      activity: 'idle',
      canSubmit: false,
      outcome: 'aborted',
      error: 'The operation stopped. Start a new conversation.',
    });
    try {
      await selected?.session.stop();
    } catch {
      /* Detached work has no publication authority. */
    }
    await cleanup;
  }
  function dispose(): Promise<void> {
    if (disposal) return disposal;
    disposed = true;
    operation?.abort();
    operation = undefined;
    void detach();
    listeners.clear();
    disposal = (async () => {
      do {
        await Promise.resolve();
        await Promise.all([...cleanups]);
      } while (cleanups.size);
    })();
    return disposal;
  }
  return {
    getSnapshot: () => snapshot,
    subscribe(notify: () => void) {
      if (disposed) return () => undefined;
      listeners.add(notify);
      return () => {
        listeners.delete(notify);
      };
    },
    submit: execute,
    newConversation,
    stop,
    dispose,
  };
}

from fastapi import FastAPI
from ag_ui_langgraph import LangGraphAgent, add_langgraph_fastapi_endpoint
from .graph import graph
 
agent = LangGraphAgent(name="streaming", graph=graph)
app = FastAPI(title="cockpit-ag-ui-streaming")
add_langgraph_fastapi_endpoint(app, agent, path="/agent")
 
 
@app.get("/ok")
def ok() -> dict:
    return {"ok": True}