613 lines
17 KiB
TypeScript

import { createLogger } from "@/lib/observability/logger";
import { JsonToSseTransformStream } from "ai";
import { extractImageUrlsFromText } from "@/lib/image-url";
import type {
ResolvedChatMediaInput,
ResolvedImageRef,
ResolvedTargetAssets,
} from "@/lib/ai/context-memory/types";
import type { CanvasViewOperationPlan } from "@/types/canvas/agent-protocol";
import type { CanvasCapabilityClientCommand } from "@/lib/canvas/capabilities/client-command";
import { generateUUID } from "@/lib/utils";
const chatApiLogger = createLogger("chat.stream");
export type StreamWriter = {
write: (chunk: any) => void;
};
export function safeWriteStreamChunk(
writer: StreamWriter,
chunk: any,
label: string
): boolean {
try {
writer.write(chunk);
return true;
} catch (error) {
chatApiLogger.warn(`[Chat API] Stream write skipped (${label})`, error);
return false;
}
}
export function writeAssistantTextChunk(
writer: StreamWriter,
text: string,
label: string,
messageId = generateUUID()
) {
safeWriteStreamChunk(
writer,
{
type: "data-appendMessage",
data: JSON.stringify({
id: messageId,
role: "assistant",
parts: [{ type: "text", text }],
}),
},
label
);
return messageId;
}
export function writeStructuredDataChunk(
writer: StreamWriter,
type: string,
data: unknown,
label: string
) {
safeWriteStreamChunk(
writer,
{
type,
data: JSON.stringify(data),
},
label
);
}
function normalizeStreamIdList(value: unknown) {
return Array.isArray(value)
? value.map((id) => String(id || "").trim()).filter(Boolean)
: [];
}
function normalizeStreamMediaUrls(value: unknown) {
return Array.isArray(value)
? value.map((url) => String(url || "").trim())
: [];
}
function isValidStreamMediaUrl(url: string) {
return /^https?:\/\//i.test(url);
}
export function normalizeCanvasMediaStreamPayload(params: {
imageUrls?: unknown;
videoUrls?: unknown;
mediaItemIds?: unknown;
handledCanvasMediaItemIds?: unknown;
}) {
const rawImageUrls = normalizeStreamMediaUrls(params.imageUrls);
const rawVideoUrls = normalizeStreamMediaUrls(params.videoUrls);
const rawEntries = [
...rawImageUrls.map((url, index) => ({
type: "image" as const,
url,
originalIndex: index,
})),
...rawVideoUrls.map((url, index) => ({
type: "video" as const,
url,
originalIndex: rawImageUrls.length + index,
})),
];
const rawMediaItemIds = normalizeStreamIdList(params.mediaItemIds);
const validEntries = rawEntries.filter((entry) =>
isValidStreamMediaUrl(entry.url)
);
const idByValidIndex = new Map<number, string>();
if (rawMediaItemIds.length === rawEntries.length) {
validEntries.forEach((entry, validIndex) => {
const id = rawMediaItemIds[entry.originalIndex];
if (id) idByValidIndex.set(validIndex, id);
});
} else if (rawMediaItemIds.length === validEntries.length) {
rawMediaItemIds.forEach((id, index) => {
if (id) idByValidIndex.set(index, id);
});
}
const dedupedEntries: Array<{
type: "image" | "video";
url: string;
mediaItemId?: string;
}> = [];
const seenMediaKeys = new Set<string>();
validEntries.forEach((entry, validIndex) => {
const key = `${entry.type}:${entry.url}`;
if (seenMediaKeys.has(key)) return;
seenMediaKeys.add(key);
dedupedEntries.push({
type: entry.type,
url: entry.url,
mediaItemId: idByValidIndex.get(validIndex),
});
});
if (
dedupedEntries.length > 0 &&
dedupedEntries.some((entry) => !entry.mediaItemId) &&
rawMediaItemIds.length === dedupedEntries.length
) {
dedupedEntries.forEach((entry, index) => {
entry.mediaItemId = rawMediaItemIds[index];
});
}
const mediaItemIds = dedupedEntries.every((entry) => entry.mediaItemId)
? dedupedEntries.map((entry) => entry.mediaItemId!)
: [];
return {
imageUrls: dedupedEntries
.filter((entry) => entry.type === "image")
.map((entry) => entry.url),
videoUrls: dedupedEntries
.filter((entry) => entry.type === "video")
.map((entry) => entry.url),
mediaItemIds,
handledCanvasMediaItemIds: normalizeStreamIdList(
params.handledCanvasMediaItemIds
),
};
}
export function writeProtocolDebugChunk(
writer: StreamWriter,
enabled: boolean,
data: unknown,
label: string
) {
if (!enabled) return;
writeStructuredDataChunk(writer, "data-protocolDebug", data, label);
}
function normalizeImageLoadingProviderType(providerType?: string | null): string {
const normalized = String(providerType || "").trim().toLowerCase();
if (normalized === "gemini-image" || normalized === "huanke-gemini") {
return "gemini";
}
if (
normalized === "alibaba-image" ||
normalized === "bfl-image" ||
normalized === "ruxa-image" ||
normalized === "fal-image" ||
normalized === "snapgen-image" ||
normalized === "openai-image" ||
normalized === "bananaimage" ||
normalized === "visionary" ||
normalized === "atlascloud" ||
normalized === "vip98k" ||
normalized === "stepfun-image" ||
normalized === "midjourney" ||
normalized === "runway" ||
normalized === "krea" ||
normalized === "volcano-vision" ||
normalized === "gemini"
) {
return normalized;
}
return "alibaba-image";
}
export function writeImageGenerationStatusChunk(
writer: StreamWriter,
params: {
phase: "queued" | "processing" | "finalizing" | "failed" | "completed";
providerType?: string | null;
modelName?: string | null;
imageCount?: number | null;
queuePosition?: number | null;
sameModelRunningCount?: number | null;
concurrencyLimit?: number | null;
estimatedWaitMs?: number | null;
},
label: string
) {
writeStructuredDataChunk(
writer,
"data-imageGenerationStatus",
{
phase: params.phase,
providerType: normalizeImageLoadingProviderType(params.providerType),
modelName: String(params.modelName || "").trim() || undefined,
imageCount:
Number.isFinite(Number(params.imageCount)) && Number(params.imageCount) > 0
? Number(params.imageCount)
: undefined,
queuePosition:
Number.isFinite(Number(params.queuePosition)) &&
Number(params.queuePosition) >= 0
? Number(params.queuePosition)
: undefined,
sameModelRunningCount:
Number.isFinite(Number(params.sameModelRunningCount)) &&
Number(params.sameModelRunningCount) >= 0
? Number(params.sameModelRunningCount)
: undefined,
concurrencyLimit:
Number.isFinite(Number(params.concurrencyLimit)) &&
Number(params.concurrencyLimit) > 0
? Number(params.concurrencyLimit)
: undefined,
estimatedWaitMs:
Number.isFinite(Number(params.estimatedWaitMs)) &&
Number(params.estimatedWaitMs) >= 0
? Number(params.estimatedWaitMs)
: undefined,
},
label
);
}
export type AgentProgressStreamPayload = {
stage: "planning" | "clarify" | "executing" | "completed" | "failed";
currentStep:
| "lock_target"
| "analyze"
| "select_tool"
| "execute"
| "finalize";
title?: string;
message?: string;
toolName?: string;
capability?: string;
};
export function writeAgentProgressChunk(
writer: StreamWriter,
payload: AgentProgressStreamPayload,
label: string
) {
writeStructuredDataChunk(
writer,
"data-agentProgress",
{
stage: payload.stage,
currentStep: payload.currentStep,
title: String(payload.title || "").trim() || undefined,
message: String(payload.message || "").trim() || undefined,
toolName: String(payload.toolName || "").trim() || undefined,
capability: String(payload.capability || "").trim() || undefined,
},
label
);
}
export function writeCanvasMediaChunk(
writer: StreamWriter,
params: {
messageId: string;
text?: string;
imageUrls?: string[];
videoUrls?: string[];
prompt?: string | null;
sourceImageUrls?: string[] | null;
agentPresetId?: string | null;
agentPresetName?: string | null;
agentCategoryId?: string | null;
agentCategoryName?: string | null;
agentExecutionSnapshotId?: string | null;
mediaItemIds?: string[] | null;
handledCanvasMediaItemIds?: string[] | null;
},
label: string
) {
const inputImageUrls = Array.isArray(params.imageUrls)
? params.imageUrls
: extractImageUrlsFromText(String(params.text || ""));
const normalizedPayload = normalizeCanvasMediaStreamPayload({
imageUrls: inputImageUrls,
videoUrls: params.videoUrls,
mediaItemIds: params.mediaItemIds,
handledCanvasMediaItemIds: params.handledCanvasMediaItemIds,
});
if (
normalizedPayload.imageUrls.length === 0 &&
normalizedPayload.videoUrls.length === 0
) {
return;
}
writeStructuredDataChunk(
writer,
"data-canvasMedia",
{
messageId: params.messageId,
imageUrls: normalizedPayload.imageUrls,
videoUrls: normalizedPayload.videoUrls,
prompt: String(params.prompt || "").trim() || undefined,
sourceImageUrls: Array.isArray(params.sourceImageUrls)
? params.sourceImageUrls
.map((url) => String(url || "").trim())
.filter((url) => /^https?:\/\//i.test(url))
: [],
agentPresetId: String(params.agentPresetId || "").trim() || undefined,
agentPresetName: String(params.agentPresetName || "").trim() || undefined,
agentCategoryId: String(params.agentCategoryId || "").trim() || undefined,
agentCategoryName: String(params.agentCategoryName || "").trim() || undefined,
agentExecutionSnapshotId:
String(params.agentExecutionSnapshotId || "").trim() || undefined,
mediaItemIds: normalizedPayload.mediaItemIds,
handledCanvasMediaItemIds: normalizedPayload.handledCanvasMediaItemIds,
},
label
);
}
export function writeCanvasOperationChunk(
writer: StreamWriter,
params: {
projectId: string;
operationId: string;
operation: string;
status?: "applied" | "failed" | "shadowed";
createdNodeIds: string[];
createdNodeIdsByMediaItemId?: Record<string, string>;
updatedNodeIds?: string[];
createdEdgeIds?: string[];
canvasRevision?: string;
anchorNodeId?: string | null;
handledCanvasMediaItemIds?: string[];
mediaItemIds?: string[];
errorCode?: string;
errorMessage?: string;
fallbackToCanvasMedia?: boolean;
message?: string;
},
label: string
) {
writeStructuredDataChunk(
writer,
"data-canvasOperation",
{
projectId: params.projectId,
operationId: params.operationId,
operation: params.operation,
status: params.status || "applied",
createdNodeIds: params.createdNodeIds || [],
createdNodeIdsByMediaItemId: params.createdNodeIdsByMediaItemId || {},
updatedNodeIds: params.updatedNodeIds || [],
createdEdgeIds: params.createdEdgeIds || [],
canvasRevision: params.canvasRevision,
anchorNodeId: params.anchorNodeId || undefined,
handledCanvasMediaItemIds:
params.status && params.status !== "applied"
? []
: normalizeStreamIdList(params.handledCanvasMediaItemIds),
mediaItemIds: normalizeStreamIdList(params.mediaItemIds),
errorCode: params.errorCode,
errorMessage: params.errorMessage,
fallbackToCanvasMedia: params.fallbackToCanvasMedia,
message: params.message,
},
label
);
}
export function writeCanvasViewOperationChunk(
writer: StreamWriter,
params: CanvasViewOperationPlan,
label: string
) {
writeStructuredDataChunk(
writer,
"data-canvasViewOperation",
{
version: 1,
viewOperationId: params.viewOperationId,
operation: params.operation,
projectId: params.projectId,
nodeIds: params.nodeIds || [],
anchorNodeId: params.anchorNodeId,
zoom: params.zoom,
padding: params.padding,
durationMs: params.durationMs,
persist: false,
center: params.center,
reason: params.reason,
},
label
);
}
export function writeCanvasCapabilityCommandChunk(
writer: StreamWriter,
command: CanvasCapabilityClientCommand,
label: string
) {
writeStructuredDataChunk(
writer,
"data-canvasCapabilityCommand",
command,
label
);
}
type AgentTargetLockRefPreview = {
url: string;
title?: string;
nodeId?: string;
assetId?: string;
source: ResolvedImageRef["source"];
role: ResolvedImageRef["role"];
isSelected?: boolean;
positionLabel?: string;
};
export type AgentTargetLockStreamPayload = {
stage: "locked" | "clarify" | "unresolved";
resolutionMode: ResolvedTargetAssets["resolutionMode"];
confidence: number;
reason: string;
needsClarification: boolean;
clarificationMessage?: string;
targetCount: number;
candidateCount: number;
targets: AgentTargetLockRefPreview[];
candidates: AgentTargetLockRefPreview[];
};
function sortTargetLockRefsForCanvasOrdering(refs: ResolvedImageRef[]) {
return [...refs].sort((left, right) => {
const leftY = Number(left.y ?? 0);
const rightY = Number(right.y ?? 0);
if (Math.abs(leftY - rightY) > 12) {
return leftY - rightY;
}
return Number(left.x ?? 0) - Number(right.x ?? 0);
});
}
function getTargetLockPositionLabel(
ref: ResolvedImageRef,
refs: ResolvedImageRef[]
): string {
if (refs.length <= 1) return "中间";
const xs = refs.map((item) => Number(item.x ?? 0));
const ys = refs.map((item) => Number(item.y ?? 0));
const minX = Math.min(...xs);
const maxX = Math.max(...xs);
const minY = Math.min(...ys);
const maxY = Math.max(...ys);
const spanX = Math.max(1, maxX - minX);
const spanY = Math.max(1, maxY - minY);
const relX = (Number(ref.x ?? 0) - minX) / spanX;
const relY = (Number(ref.y ?? 0) - minY) / spanY;
const horizontal = relX <= 0.33 ? "左" : relX >= 0.66 ? "右" : "中";
const vertical = relY <= 0.33 ? "上" : relY >= 0.66 ? "下" : "中";
if (horizontal === "中" && vertical === "中") return "中间";
return `${vertical}${horizontal}`;
}
function serializeAgentTargetLockRefs(refs: ResolvedImageRef[]) {
const orderedRefs = sortTargetLockRefsForCanvasOrdering(refs).slice(0, 4);
return orderedRefs.map((ref) => ({
url: ref.url,
title: String(ref.title || "").trim() || undefined,
nodeId: ref.nodeId,
assetId: ref.assetId,
source: ref.source,
role: ref.role,
isSelected: ref.isSelected === true,
positionLabel: getTargetLockPositionLabel(ref, orderedRefs),
})) satisfies AgentTargetLockRefPreview[];
}
export function buildAgentTargetLockPayload(params: {
requestTargetResolution: ResolvedTargetAssets;
requestResolvedMediaInput: ResolvedChatMediaInput;
needsClarification: boolean;
clarificationMessage?: string;
}): AgentTargetLockStreamPayload {
const stage = params.needsClarification
? "clarify"
: params.requestTargetResolution.targetRefs.length > 0
? "locked"
: "unresolved";
return {
stage,
resolutionMode: params.requestTargetResolution.resolutionMode,
confidence: Number(params.requestTargetResolution.confidence || 0),
reason: params.requestTargetResolution.reason,
needsClarification: params.needsClarification,
clarificationMessage: params.clarificationMessage?.trim() || undefined,
targetCount: params.requestTargetResolution.targetRefs.length,
candidateCount: params.requestTargetResolution.candidateRefs.length,
targets: serializeAgentTargetLockRefs(params.requestTargetResolution.targetRefs),
candidates: serializeAgentTargetLockRefs(
params.requestTargetResolution.candidateRefs.length > 0
? params.requestTargetResolution.candidateRefs
: params.requestResolvedMediaInput.imageRefs
),
};
}
export function writeAgentTargetLockChunk(
writer: StreamWriter,
payload: AgentTargetLockStreamPayload,
label: string
) {
writeStructuredDataChunk(writer, "data-agentTargetLock", payload, label);
}
export function createSseResponse(stream: ReadableStream) {
return new Response(stream.pipeThrough(new JsonToSseTransformStream()), {
headers: {
"Content-Type": "text/event-stream; charset=utf-8",
"Cache-Control": "no-cache, no-transform",
Connection: "keep-alive",
"X-Accel-Buffering": "no",
},
});
}
export function startStreamKeepalive(
writer: StreamWriter,
label: string,
intervalMs = 15_000
) {
const timer = setInterval(() => {
safeWriteStreamChunk(
writer,
{
type: "data-keepalive",
data: JSON.stringify({ ts: Date.now(), label }),
},
`${label}:keepalive`
);
}, intervalMs);
return () => clearInterval(timer);
}
export async function runBoundedParallel<T>(
items: T[],
workerCount: number,
worker: (item: T, index: number) => Promise<void>
) {
const maxWorkers = Math.max(1, Math.min(workerCount, items.length));
let cursor = 0;
let firstError: unknown = null;
const workers = Array.from({ length: maxWorkers }, async () => {
while (cursor < items.length) {
const index = cursor;
cursor += 1;
if (firstError) {
return;
}
try {
await worker(items[index]!, index);
} catch (error) {
firstError = firstError || error;
}
}
});
await Promise.all(workers);
if (firstError) {
throw firstError;
}
}