Jyotisha/frontend/src/lib/stream-agent-response.ts

import {
  appendConsultationRuntimeStep,
  type ConsultationRuntimeState,
} from "../mastra/consultation-tools.ts";
import {
  agentExecutionReceiptSchema,
  consultationAgentPublicEventSchema,
  publicActivityPhaseSchema,
  type AgentExecutionReceipt,
  type ConsultationAgentPublicEvent,
} from "./consultation-agent-events.ts";
import { toAgentModelFinishReason } from "./agent-observability.ts";
import { createVisibleTextTransformer } from "./stream-text-response.ts";
import { sanitizePublicThinkingText } from "./public-thinking.ts";

type Chunk = { type?: string; payload?: Record<string, unknown>; data?: unknown };
type ChunkStream = AsyncIterable<unknown> | ReadableStream<unknown>;

async function* readChunks(stream: ChunkStream): AsyncIterable<Chunk> {
  const values = Symbol.asyncIterator in stream
    ? stream as AsyncIterable<unknown>
    : (async function* () {
      const reader = (stream as ReadableStream<unknown>).getReader();
      try {
        while (true) {
          const { done, value } = await reader.read();
          if (done) return;
          yield value;
        }
      } finally {
        reader.releaseLock();
      }
    })();
  for await (const value of values) {
    if (value && typeof value === "object") yield value as Chunk;
  }
}
type Status = "ready" | "degraded" | "blocked";

type EventOptions = {
  runId: string;
  requestId: string;
  toolStatus: () => Status;
  receipt: () => AgentExecutionReceipt;
  state?: ConsultationRuntimeState;
};

function activity(value: unknown): ConsultationAgentPublicEvent | null {
  if (!value || typeof value !== "object") return null;
  const data = value as { phase?: unknown; label?: unknown };
  const phase = publicActivityPhaseSchema.safeParse(data.phase);
  if (!phase.success || typeof data.label !== "string") return null;
  return { type: "activity", phase: phase.data, label: data.label.slice(0, 120) };
}

/**
 * The runtime only reveals how the model loop ended through the stream: one
 * `step-finish` per model step, then a terminal `finish` carrying the reason
 * the model stopped and the authoritative step list. Without this, a run that
 * exhausted its step budget is indistinguishable from one that chose to stop,
 * because progressive-disclosure reads never reach the public event stream.
 */
function finishTelemetry(chunk: Chunk) {
  const payload = chunk.payload as {
    stepResult?: { reason?: unknown };
    output?: { steps?: unknown };
  } | undefined;
  const steps = payload?.output?.steps;
  return {
    reason: toAgentModelFinishReason(payload?.stepResult?.reason),
    stepCount: Array.isArray(steps) ? steps.length : null,
  };
}

function safeToolError(error: unknown) {
  if (error instanceof DOMException && error.name === "AbortError") return "cancelled" as const;
  if (error instanceof DOMException && error.name === "TimeoutError") return "timeout" as const;
  return "calculation_failed" as const;
}

function isTimeoutOrAbort(error: unknown) {
  return error instanceof Error && (error.name === "TimeoutError" || error.name === "AbortError");
}

type RunFailedCode = "runtime_contract_incomplete" | "empty_answer" | "answer_truncated" | "calculation_failed";

function runFailedCode(error: unknown, emitted: boolean): RunFailedCode {
  if (error instanceof Error && error.message === "runtime_contract_incomplete") return "runtime_contract_incomplete";
  if (error instanceof Error && error.message === "empty_answer") return "empty_answer";
  if (error instanceof Error && error.message === "answer_truncated") return "answer_truncated";
  if (emitted && isTimeoutOrAbort(error)) return "answer_truncated";
  return "calculation_failed";
}

function runFailedMessage(code: RunFailedCode) {
  if (code === "runtime_contract_incomplete") return "Agent 未完成必要的方法与计算步骤，本次不会扣点。";
  if (code === "empty_answer") return "计算已完成，但这次没有生成回答，本次不会扣点。请再发送一次。";
  if (code === "answer_truncated") return "回答未完成，已保留现有内容；本次不会扣点。";
  return "咨询暂时无法完成，本次不会扣点。";
}

/**
 * Whether a tool result is really an input rejection Mastra resolved with.
 *
 * Mastra validates arguments against the tool's inputSchema before `execute`
 * and reports a mismatch by *resolving* with an error envelope rather than
 * throwing. Passing that through as `tool.completed` would tell the client a
 * calculation finished while the tool body never ran. Matched on the envelope's
 * shape, not its English message, so an upstream wording change cannot silently
 * turn a rejection back into a success.
 */
function isToolInputRejection(result: unknown) {
  if (!result || typeof result !== "object") return false;
  const envelope = result as { error?: unknown; validationErrors?: unknown };
  return envelope.error === true
    && typeof envelope.validationErrors === "object"
    && envelope.validationErrors !== null;
}

/**
 * Record a tool failure the tool itself could not record.
 *
 * A call rejected against the tool's input schema never enters `execute`, so
 * nothing in the tool runs: the run reported a `tool.failed` event to the client
 * while its receipt showed no failed step and its budget counted no call. The
 * stream is the only place that observes every failure, whichever side of the
 * schema it came from.
 *
 * The tool still records the failures it can see, with the duration and cause it
 * alone knows, so this only fills the gap: it appends when the stream has seen
 * more tool errors than the state has failed tool steps. The error object here
 * is provider-shaped and cannot be classified further, but the call not reaching
 * `execute` is itself the diagnosis.
 */
function recordUnrecordedToolFailure(
  state: ConsultationRuntimeState,
  toolErrorsSeen: number,
  durationMs: number,
) {
  const recorded = state.steps.filter((step) => step.kind === "tool" && step.status === "failed").length;
  if (recorded >= toolErrorsSeen) return;
  appendConsultationRuntimeStep(state, {
    kind: "tool",
    name: "run-jyotish-consultation",
    status: "failed",
    durationMs,
    failureCode: "tool_call_rejected",
  });
}

/**
 * The two events that used to report the model activating the skill. The server
 * binds the method into the instructions now, so they are announced once at the
 * start of a run instead of being read off a tool call that no longer happens.
 * They stay in the public stream because they are what tells a waiting client
 * that method is in hand, and the first of them is a run's first activity.
 */
const skillBoundEvents: readonly ConsultationAgentPublicEvent[] = [
  { type: "skill.started", name: "jyotish-vedic-astrology" },
  { type: "skill.completed", name: "jyotish-vedic-astrology" },
];

function mapChunk(
  chunk: Chunk,
  options: EventOptions,
  startedAt: Map<string, number>,
  toolErrors: { seen: number },
): ConsultationAgentPublicEvent[] {
  const payload = chunk.payload ?? {};
  if (chunk.type === "data-jyotish-activity") {
    const event = activity(chunk.data);
    return event ? [event] : [];
  }
  if (chunk.type === "tool-call") {
    const toolName = payload.toolName;
    const callId = typeof payload.toolCallId === "string" ? payload.toolCallId : "tool";
    if (toolName === "run-jyotish-consultation") {
      startedAt.set(callId, Date.now());
      return [{ type: "tool.started", callId, tool: "run-jyotish-consultation", label: "正在计算个人星盘" }];
    }
  }
  if (chunk.type === "tool-result") {
    const toolName = payload.toolName;
    const callId = typeof payload.toolCallId === "string" ? payload.toolCallId : "tool";
    if (toolName === "run-jyotish-consultation") {
      const durationMs = Math.max(0, Date.now() - (startedAt.get(callId) ?? Date.now()));
      if (isToolInputRejection(payload.result)) {
        toolErrors.seen += 1;
        if (options.state) recordUnrecordedToolFailure(options.state, toolErrors.seen, durationMs);
        return [{ type: "tool.failed", callId, tool: "run-jyotish-consultation", code: "calculation_failed" }];
      }
      return [{
        type: "tool.completed", callId, tool: "run-jyotish-consultation", status: options.toolStatus(),
        durationMs,
      }];
    }
  }
  if (chunk.type === "tool-error") {
    const toolName = payload.toolName;
    if (toolName === "run-jyotish-consultation") {
      const callId = typeof payload.toolCallId === "string" ? payload.toolCallId : "tool";
      toolErrors.seen += 1;
      if (options.state) {
        recordUnrecordedToolFailure(
          options.state,
          toolErrors.seen,
          Math.max(0, Date.now() - (startedAt.get(callId) ?? Date.now())),
        );
      }
      return [{
        type: "tool.failed",
        callId,
        tool: "run-jyotish-consultation",
        code: safeToolError(payload.error),
      }];
    }
  }
  return [];
}

export async function collectAgentPublicEvents(stream: ChunkStream | Iterable<Chunk>, options: EventOptions) {
  const events: ConsultationAgentPublicEvent[] = [
    { type: "run.started", runId: options.runId, requestId: options.requestId },
    ...skillBoundEvents,
  ];
  const startedAt = new Map<string, number>();
  const toolErrors = { seen: 0 };
  for await (const chunk of stream instanceof ReadableStream || Symbol.asyncIterator in stream ? readChunks(stream as ChunkStream) : stream) {
    events.push(...mapChunk(chunk, options, startedAt, toolErrors));
    if (chunk.type === "text-delta" && typeof chunk.payload?.text === "string") {
      events.push({ type: "answer.delta", text: chunk.payload.text });
    }
    if (chunk.type === "reasoning-delta" && typeof chunk.payload?.text === "string") {
      const thinking = sanitizePublicThinkingText(chunk.payload.text);
      if (thinking) events.push({ type: "thinking.delta", text: thinking });
    }
  }
  events.push({ type: "run.completed", receipt: agentExecutionReceiptSchema.parse(options.receipt()) });
  return events.map((event) => consultationAgentPublicEventSchema.parse(event));
}

type StreamAgentResponseOptions = EventOptions & {
  state: ConsultationRuntimeState;
  stream: ChunkStream;
  transformText?: (text: string) => string;
  requireTool: boolean;
  retry?: () => Promise<ChunkStream>;
  retryForAnswer?: () => Promise<ChunkStream>;
  continueAfterDisconnect?: boolean;
  headers?: HeadersInit;
  onFirstActivity?: () => void | Promise<void>;
  onFirstOutput?: () => void | Promise<void>;
  onComplete?: (output: string, receipt: AgentExecutionReceipt, thinkingText?: string) => void | Promise<void>;
  onError?: (error: unknown, emitted: boolean, output: string) => void | Promise<void>;
  onCancel?: (emitted: boolean) => void | Promise<void>;
};

// Failed attempts are retried by the model against the same request-scoped
// calculation cache, so only successful workflow executions may count against
// the single-calculation boundary. Gating on total attempts would make any
// transient failure permanently unrecoverable.
function contractReady(options: StreamAgentResponseOptions) {
  return options.state.jyotishSkillBound
    && (!options.requireTool || (options.state.consultationToolCompleted && options.state.consultationToolSuccessCount === 1));
}

export function streamAgentResponse(options: StreamAgentResponseOptions) {
  const encoder = new TextEncoder();
  let disconnected = false;
  let settled = false;
  let settling = false;
  let emitted = false;
  let firstActivity = false;
  let firstOutput = false;
  let fullOutput = "";
  let fullThinking = "";
  const startedAt = new Map<string, number>();
  // A retry reuses these counters so a failure in either attempt is recorded once.
  const toolErrors = { seen: 0 };
  const send = (controller: ReadableStreamDefaultController<Uint8Array> | undefined, event: ConsultationAgentPublicEvent) => {
    if (!firstActivity && (event.type === "skill.started" || event.type === "tool.started" || event.type === "activity")) {
      firstActivity = true;
      void Promise.resolve(options.onFirstActivity?.()).catch(() => {});
    }
    if (!disconnected && controller) controller.enqueue(encoder.encode(`${JSON.stringify(consultationAgentPublicEventSchema.parse(event))}\n`));
  };

  async function consumeAttempt(controller: ReadableStreamDefaultController<Uint8Array> | undefined, stream: ChunkStream) {
    const visible = createVisibleTextTransformer(options.transformText ?? ((value) => value));
    let held = "";
    let composingSent = false;
    const outputText = async (text: string) => {
      // Text the model writes before the contract is ready is not the answer: it
      // is the model narrating its own in-progress or failed tool calls. Holding
      // it meant a later successful call released that narration as the entire
      // visible answer, so a run where the model recovered read as a run where it
      // explained itself instead of answering. Drop it.
      if (!contractReady(options)) return;
      held += text;
      if (!held) return;
      if (!composingSent) {
        composingSent = true;
        send(controller, { type: "activity", phase: "answer-composition", label: "正在组织回答" });
      }
      if (!firstOutput && /\S/.test(held)) {
        firstOutput = true;
        await options.onFirstOutput?.();
      }
      send(controller, { type: "answer.delta", text: held });
      fullOutput += held;
      if (/\S/.test(held)) emitted = true;
      held = "";
    };
    // A retry runs a second model loop under the same step budget, so the run
    // total accumulates while the finish reason describes the latest attempt.
    const stepCountBeforeAttempt = options.state.modelStepCount;
    try {
      for await (const chunk of readChunks(stream)) {
        for (const event of mapChunk(chunk, options, startedAt, toolErrors)) send(controller, event);
        if (chunk.type === "step-finish") options.state.modelStepCount += 1;
        if (chunk.type === "finish") {
          const finish = finishTelemetry(chunk);
          options.state.modelFinishReason = finish.reason;
          if (finish.stepCount !== null) options.state.modelStepCount = stepCountBeforeAttempt + finish.stepCount;
        }
        if (chunk.type === "text-delta" && typeof chunk.payload?.text === "string") {
          await outputText(visible.push(chunk.payload.text));
        }
        if (chunk.type === "reasoning-delta" && typeof chunk.payload?.text === "string") {
          const thinking = sanitizePublicThinkingText(chunk.payload.text);
          if (thinking) {
            fullThinking = `${fullThinking}${thinking}`.slice(0, 4_000);
            send(controller, { type: "thinking.delta", text: thinking });
          }
        }
      }
      await outputText(visible.finish(""));
    } catch (error) {
      try {
        await outputText(visible.finish(""));
      } catch {}
      throw error;
    }
  }

  const body = new ReadableStream<Uint8Array>({
    start(controller) {
      void (async () => {
        send(controller, { type: "run.started", runId: options.runId, requestId: options.requestId });
        for (const event of skillBoundEvents) send(controller, event);
        try {
          await consumeAttempt(controller, options.stream);
          if (!contractReady(options) && options.retry) {
            appendConsultationRuntimeStep(options.state, { kind: "validation", name: "runtime-contract-retry", status: "completed" });
            send(controller, { type: "activity", phase: "loading-method", label: "正在补齐方法与计算步骤" });
            await consumeAttempt(controller, await options.retry());
          }
          if (!contractReady(options)) throw new Error("runtime_contract_incomplete");
          // A run whose calculation succeeded and whose model then wrote nothing
          // used to be answered with a fixed apology and billed as a completed
          // consultation: the user paid for a sentence saying there was nothing
          // to say. Ask once more instead. The calculation is cached for the
          // request, so the second attempt re-reads the same evidence without
          // recomputing it, and whatever ended the first attempt—an exhausted
          // step budget above all—does not carry into a fresh model loop.
          if (!/\S/.test(fullOutput) && options.retryForAnswer) {
            appendConsultationRuntimeStep(options.state, { kind: "validation", name: "answer-retry", status: "completed" });
            send(controller, { type: "activity", phase: "answer-composition", label: "正在组织回答" });
            await consumeAttempt(controller, await options.retryForAnswer());
          }
          // Still nothing to show. Failing is the honest outcome and it is the
          // one that does not charge for the run.
          if (!/\S/.test(fullOutput)) throw new Error("empty_answer");
          // A spoken answer that stopped because the token budget ran out is
          // not a completed consultation. The heading may already be on screen,
          // so keep it and refuse to bill.
          if (options.state.modelFinishReason === "length") throw new Error("answer_truncated");
          settling = true;
          const receipt = agentExecutionReceiptSchema.parse(options.receipt());
          await options.onComplete?.(fullOutput, receipt, fullThinking || undefined);
          settled = true;
          settling = false;
          send(controller, { type: "run.completed", receipt });
          if (!disconnected) controller.close();
        } catch (error) {
          if (settled) return;
          settled = true;
          settling = false;
          try {
            await options.onError?.(error, emitted, fullOutput);
          } catch {}
          const code = runFailedCode(error, emitted);
          // Step durations, the step budget and the workflow route are the only
          // evidence the caller has for why a run failed. Building the receipt
          // must not be able to replace the failure event with a silent close.
          let failureReceipt: AgentExecutionReceipt | undefined;
          try {
            failureReceipt = agentExecutionReceiptSchema.parse(options.receipt());
          } catch {}
          send(controller, {
            type: "run.failed",
            code,
            message: runFailedMessage(code),
            ...(failureReceipt ? { receipt: failureReceipt } : {}),
          });
          if (!disconnected) controller.close();
        }
      })();
    },
    async cancel() {
      if (settled) return;
      if (settling || options.continueAfterDisconnect) {
        disconnected = true;
        return;
      }
      settled = true;
      await options.onCancel?.(emitted);
    },
  });
  return new Response(body, {
    headers: {
      "cache-control": "no-cache, no-transform",
      "content-type": "application/x-ndjson; charset=utf-8",
      "x-accel-buffering": "no",
      "x-ayanam-mode": "mastra-agentic",
      "x-ayanam-request-id": options.requestId,
      ...options.headers,
    },
  });
}