import type { StreamFn } from "@earendil-works/pi-agent-core";

import { createAssistantMessageEventStream, streamSimple } from "@earendil-works/pi-ai";

import { streamSimple } from "@earendil-works/pi-ai";

import { createSubsystemLogger } from "openclaw/plugin-sdk/logging-core";

import type { ProviderWrapStreamFnContext } from "openclaw/plugin-sdk/plugin-entry";

import { createPlainTextToolCallPromotionWrapper } from "openclaw/plugin-sdk/provider-stream-shared";

import { ssrfPolicyFromHttpBaseUrlAllowedHostname } from "openclaw/plugin-sdk/ssrf-runtime";

import { LMSTUDIO_PROVIDER_ID } from "./defaults.js";

import { ensureLmstudioModelLoaded } from "./models.fetch.js";

import { resolveLmstudioInferenceBase } from "./models.js";

import {

createLmstudioSyntheticToolCallId,

parseLmstudioPlainTextToolCalls,

} from "./plain-text-tool-calls.js";

import { resolveLmstudioProviderHeaders, resolveLmstudioRuntimeApiKey } from "./runtime.js";

const log = createSubsystemLogger("extensions/lmstudio/stream");

type StreamOptions = Parameters<StreamFn>[2];

type StreamModel = Parameters<StreamFn>[0];

type StreamContext = Parameters<StreamFn>[1];

const preloadInFlight = new Map<string, Promise<void>>();

/**

};

}

function resolveContextToolNames(context: StreamContext): Set<string> {

const tools = (context as { tools?: unknown }).tools;

if (!Array.isArray(tools)) {

return new Set();

}

const names = tools

.map((tool) => {

const record = toRecord(tool);

return typeof record?.name === "string" && record.name.trim() ? record.name : undefined;

})

.filter((name): name is string => Boolean(name));

return new Set(names);

}

function couldStillBePlainTextToolCall(text: string): boolean {

if (text.length > 256_000) {

return false;

}

const trimmed = text.trimStart();

return (

trimmed.length === 0 ||

trimmed.startsWith("[") ||

trimmed.startsWith("<|channel|>") ||

trimmed.startsWith("commentary") ||

trimmed.startsWith("analysis") ||

trimmed.startsWith("final")

);

}

function createLmstudioToolCallBlock(parsed: {

arguments: Record<string, unknown>;

name: string;

}): Record<string, unknown> {

return {

type: "toolCall",

id: createLmstudioSyntheticToolCallId(),

name: parsed.name,

arguments: parsed.arguments,

partialArgs: JSON.stringify(parsed.arguments),

};

}

function promoteLmstudioPlainTextToolCalls(

message: unknown,

toolNames: Set<string>,

): Record<string, unknown> | undefined {

const messageRecord = toRecord(message);

if (!messageRecord) {

return undefined;

}

if (!Array.isArray(messageRecord.content)) {

if (typeof messageRecord.content !== "string" || !messageRecord.content.trim()) {

return undefined;

}

const parsed = parseLmstudioPlainTextToolCalls(messageRecord.content, toolNames);

if (!parsed) {

return undefined;

}

return {

...messageRecord,

content: parsed.map(createLmstudioToolCallBlock),

stopReason: "toolUse",

};

}

if (

messageRecord.content.some((block) => toRecord(block)?.type === "toolCall") ||

messageRecord.content.length === 0

) {

return undefined;

}

let promoted = false;

const nextContent: Array<Record<string, unknown>> = [];

for (const block of messageRecord.content) {

const blockRecord = toRecord(block);

if (!blockRecord) {

return undefined;

}

if (blockRecord.type !== "text") {

nextContent.push(blockRecord);

continue;

}

const text = typeof blockRecord.text === "string" ? blockRecord.text : "";

if (!text.trim()) {

continue;

}

const parsed = parseLmstudioPlainTextToolCalls(text, toolNames);

if (!parsed) {

return undefined;

}

nextContent.push(...parsed.map(createLmstudioToolCallBlock));

promoted = true;

}

if (!promoted) {

return undefined;

}

return {

...messageRecord,

content: nextContent,

stopReason: "toolUse",

};

}

function emitPromotedToolCallEvents(

stream: { push(event: unknown): void },

message: Record<string, unknown>,

): void {

const content = Array.isArray(message.content) ? message.content : [];

content.forEach((block, contentIndex) => {

const record = toRecord(block);

if (record?.type !== "toolCall") {

return;

}

stream.push({ type: "toolcall_start", contentIndex, partial: message });

stream.push({

type: "toolcall_delta",

contentIndex,

delta: typeof record.partialArgs === "string" ? record.partialArgs : "{}",

partial: message,

});

});

}

function wrapLmstudioPlainTextToolCalls(

source: ReturnType<StreamFn>,

context: StreamContext,

): ReturnType<StreamFn> {

const toolNames = resolveContextToolNames(context);

if (toolNames.size === 0) {

return source;

}

const output = createAssistantMessageEventStream();

const stream = output as unknown as { push(event: unknown): void; end(): void };

void (async () => {

const bufferedTextEvents: unknown[] = [];

let bufferedText = "";

let ended = false;

const endStream = () => {

if (!ended) {

ended = true;

stream.end();

}

};

const flushBufferedTextEvents = () => {

for (const event of bufferedTextEvents.splice(0)) {

stream.push(event);

}

bufferedText = "";

};

try {

for await (const event of source as AsyncIterable<unknown>) {

const record = toRecord(event);

const type = typeof record?.type === "string" ? record.type : "";

if (type === "text_start" || type === "text_delta" || type === "text_end") {

bufferedTextEvents.push(event);

if (typeof record?.delta === "string") {

bufferedText += record.delta;

} else if (typeof record?.content === "string" && !bufferedText) {

bufferedText = record.content;

}

if (!couldStillBePlainTextToolCall(bufferedText)) {

flushBufferedTextEvents();

}

continue;

}

if (type === "done") {

const promotedMessage = promoteLmstudioPlainTextToolCalls(record?.message, toolNames);

if (promotedMessage) {

bufferedTextEvents.splice(0);

bufferedText = "";

emitPromotedToolCallEvents(stream, promotedMessage);

stream.push({ ...record, reason: "toolUse", message: promotedMessage });

} else {

flushBufferedTextEvents();

stream.push(event);

}

endStream();

return;

}

flushBufferedTextEvents();

stream.push(event);

if (type === "error") {

endStream();

return;

}

}

flushBufferedTextEvents();

} catch (error) {

stream.push({

type: "error",

reason: "error",

error: {

role: "assistant",

content: [],

stopReason: "error",

errorMessage: error instanceof Error ? error.message : String(error),

},

});

} finally {

endStream();

}

})();

return output as ReturnType<StreamFn>;

}

function createPreloadKey(params: {

baseUrl: string;

modelKey: string;

export function wrapLmstudioInferencePreload(ctx: ProviderWrapStreamFnContext): StreamFn {

const underlying = ctx.streamFn ?? streamSimple;

const streamWithPlainTextToolCalls = createPlainTextToolCallPromotionWrapper(underlying);

return (model, context, options) => {

if (model.provider !== LMSTUDIO_PROVIDER_ID) {

return underlying(model, context, options);

}

const providerConfig = ctx.config?.models?.providers?.[LMSTUDIO_PROVIDER_ID];

if (!shouldPreloadLmstudioModels(providerConfig)) {

const stream = underlying(withLmstudioUsageCompat(model), context, options);

return (async () => {

const resolvedStream = stream instanceof Promise ? await stream : stream;

return wrapLmstudioPlainTextToolCalls(resolvedStream, context);

})();

return streamWithPlainTextToolCalls(withLmstudioUsageCompat(model), context, options);

}

const providerBaseUrl = providerConfig?.baseUrl;

const resolvedBaseUrl = resolveLmstudioInferenceBase(

// LM Studio uses OpenAI-compatible streaming usage payloads when requested via

// `stream_options.include_usage`. Force this compat flag at call time so usage

// reporting remains enabled even when catalog entries omitted compat metadata.

const stream = underlying(withLmstudioUsageCompat(model), context, options);

const stream = streamWithPlainTextToolCalls(withLmstudioUsageCompat(model), context, options);

const resolvedStream = stream instanceof Promise ? await stream : stream;

return wrapLmstudioPlainTextToolCalls(resolvedStream, context);

return resolvedStream;

})();

};

}