Files
2026-07-13 13:39:12 +08:00

482 lines
17 KiB
TypeScript

/**
* TraeExecutor — talks to Trae's remote agent API (solo_agent_remote).
*
* Flow (reverse-engineered from solo.trae.ai web client):
* 1. POST {base}/chat_sessions → { data: { chat_session_id, message_id } }
* 2. GET {base}/chat_sessions/{id}/events?reply_to_message_id={message_id}
* → text/event-stream. Assistant text streams in `plan_item` events under
* the `thought` field (cumulative per plan-item id). `token_usage` carries
* usage; `done` ends the turn; `error` carries upstream errors.
*
* Auth: header `Authorization: Cloud-IDE-JWT <JWT>` (RS256, ~14-day lifetime).
* The JWT is stored as credentials.accessToken; identity fields (web_id,
* biz_user_id, user_unique_id, scope, tenant, region) live in providerSpecificData.
*
* Model selection: model="auto" → server picks; otherwise model is the upstream
* `name` from GET {base}/models (e.g. gpt-5.2, gemini-3.1-pro, kimi-k2.5).
*/
import { BaseExecutor, mergeUpstreamExtraHeaders } from "./base.ts";
import { PROVIDERS } from "../config/constants.ts";
import { sanitizeErrorMessage } from "../utils/error.ts";
type JsonRecord = Record<string, unknown>;
type ChatMessage = { role?: string; content?: unknown };
const STREAM_TIMEOUT_MS = parseInt(process.env.TRAE_STREAM_TIMEOUT_MS || "300000", 10);
function flattenQuery(messages: ChatMessage[]): string {
const parts: string[] = [];
for (const m of messages) {
let content = "";
if (typeof m.content === "string") content = m.content;
else if (Array.isArray(m.content)) {
content = m.content
.map((p) => {
if (typeof p === "string") return p;
if (p && typeof p === "object") return String((p as JsonRecord).text ?? "");
return "";
})
.join("");
}
if (m.role === "system") parts.push(`[System]\n${content}`);
else if (m.role === "assistant") parts.push(`[Assistant]\n${content}`);
else parts.push(content);
}
const text = parts.join("\n\n");
// Trae expects query as a JSON-encoded string of typed content blocks.
return JSON.stringify([{ type: "text", data: { content: text } }]);
}
export class TraeExecutor extends BaseExecutor {
constructor() {
super("trae", PROVIDERS["trae"]);
}
private base(): string {
return (this.config.baseUrl || "https://core-normal.trae.ai/api/remote/v1").replace(/\/$/, "");
}
buildHeaders(credentials): Record<string, string> {
const token = (credentials.accessToken as string) || "";
const psd = (credentials.providerSpecificData as JsonRecord) || {};
return {
Authorization: `Cloud-IDE-JWT ${token}`,
"Content-Type": "application/json",
"X-Trae-Client-Type": "web",
"X-Preferenced-Language": (psd.appLanguage as string) || "en",
"x-user-region": (psd.userRegion as string) || "US",
Referer: "https://solo.trae.ai/",
"User-Agent":
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 " +
"(KHTML, like Gecko) Chrome/149.0.0.0 Safari/537.36",
};
}
/**
* SOLO exposes two session modes (the toggle on solo.trae.ai):
* - "code" (default): full model picker — `auto` plus named models
* (gpt-5.4, kimi-k2.5, gemini-3.1-pro, …).
* - "work": a single, faster "auto" agent with no model picker.
* We surface "work" as its own model id (`trae/work`) so callers can opt into
* the fast lane; any other model id runs in "code" mode. "work" forces the
* auto strategy with an empty model_name, since it has no model selection.
*/
private resolveMode(model: string): {
mode: "code" | "work";
strategy: "auto" | "manual";
modelName: string;
} {
const m = (model || "").trim().toLowerCase();
if (m === "work" || m === "auto-work" || m === "solo-work") {
return { mode: "work", strategy: "auto", modelName: "" };
}
const auto = !m || m === "auto";
return { mode: "code", strategy: auto ? "auto" : "manual", modelName: auto ? "" : model };
}
private commonParams(psd: JsonRecord, mode: "code" | "work", sessionId?: string): string {
const cp: JsonRecord = {
language: "en-us",
app_language: (psd.appLanguage as string) || "en",
quality: "stable",
app_version: (psd.appVersion as string) || "1.0.0.1229",
web_id: (psd.webId as string) || "",
user_identity: (psd.userIdentity as string) || "Free",
is_freshman: "0",
biz_user_id: (psd.bizUserId as string) || "",
user_unique_id: (psd.userUniqueId as string) || "",
scope: (psd.scope as string) || "marscode-us",
tenant: (psd.tenant as string) || "marscode",
region: (psd.region as string) || "US-East",
aiRegion: (psd.aiRegion as string) || (psd.region as string) || "US-East",
is_privacy_mode: 0,
privacy_mode: "off",
solo_chat_mode: mode,
};
if (sessionId) cp.biz_session_id = sessionId;
return JSON.stringify(cp);
}
/** POST /chat_sessions — creates a session and submits the first turn. */
private async createSession(
headers: Record<string, string>,
query: string,
model: string,
psd: JsonRecord,
signal?: AbortSignal | null
): Promise<{ sessionId: string; messageId: string }> {
const { mode, strategy, modelName } = this.resolveMode(model);
const body = {
mode,
environment_id: "default",
initial_message: {
chat_session_id: "",
content: [],
query,
model_name: modelName,
agent_type: "solo_agent_remote",
model_selection_strategy: strategy,
common_params: this.commonParams(psd, mode),
},
env: "remote",
auto_create_project: false,
origin: "web",
};
const res = await fetch(`${this.base()}/chat_sessions`, {
method: "POST",
headers,
body: JSON.stringify(body),
signal: signal || undefined,
});
const text = await res.text();
if (!res.ok) throw new Error(`[${res.status}] ${text}`);
const json = JSON.parse(text);
if (json?.code !== 0) throw new Error(`Trae create_session: ${JSON.stringify(json)}`);
return { sessionId: json.data.chat_session_id, messageId: json.data.message_id };
}
/**
* GET /events SSE → invoke onEvent(eventType, dataObj) per frame.
* Resolves when `done`/`error` arrives, the stream ends, or timeout fires.
*/
private async streamEvents(
headers: Record<string, string>,
sessionId: string,
replyTo: string,
onEvent: (ev: string | null, data: JsonRecord) => boolean,
signal?: AbortSignal | null
): Promise<void> {
const url = `${this.base()}/chat_sessions/${sessionId}/events?reply_to_message_id=${encodeURIComponent(replyTo)}`;
const ctrl = new AbortController();
// If the caller's signal is already aborted, abort upfront so we don't open
// a network request the consumer no longer wants.
if (signal?.aborted) ctrl.abort();
const timer = setTimeout(() => ctrl.abort(new Error("trae stream timeout")), STREAM_TIMEOUT_MS);
const onAbort = () => ctrl.abort();
if (signal) signal.addEventListener("abort", onAbort, { once: true });
try {
const res = await fetch(url, { method: "GET", headers, signal: ctrl.signal });
if (!res.ok || !res.body) throw new Error(`[${res.status}] events stream failed`);
const reader = res.body.getReader();
const decoder = new TextDecoder();
let buf = "";
let ev: string | null = null;
for (;;) {
const { done, value } = await reader.read();
if (done) break;
buf += decoder.decode(value, { stream: true });
let nl: number;
// SSE frames are separated by lines; process complete lines only.
while ((nl = buf.indexOf("\n")) >= 0) {
const line = buf.slice(0, nl).replace(/\r$/, "");
buf = buf.slice(nl + 1);
if (line.startsWith("event:")) ev = line.slice(6).trim();
else if (line.startsWith("data:")) {
const payload = line.slice(5).trim();
let data: JsonRecord;
try {
data = JSON.parse(payload);
} catch {
data = { _raw: payload };
}
if (onEvent(ev, data)) {
await reader.cancel().catch(() => {});
return;
}
} else if (line === "") ev = null;
}
}
} finally {
clearTimeout(timer);
if (signal) signal.removeEventListener("abort", onAbort);
}
}
async execute({ model, body, stream, credentials, signal, upstreamExtraHeaders }) {
const headers = this.buildHeaders(credentials as JsonRecord);
mergeUpstreamExtraHeaders(headers, upstreamExtraHeaders as Record<string, string> | null);
const psd = ((credentials as JsonRecord).providerSpecificData as JsonRecord) || {};
const reqBody = body as { messages?: ChatMessage[] };
const query = flattenQuery(reqBody.messages || []);
const responseId = `chatcmpl-trae-${Date.now()}`;
const created = Math.floor(Date.now() / 1000);
const errResponse = (status: number, message: string) =>
new Response(
JSON.stringify({
error: { message: sanitizeErrorMessage(message), type: "api_error", code: "" },
}),
{ status, headers: { "Content-Type": "application/json" } }
);
let session: { sessionId: string; messageId: string };
try {
session = await this.createSession(
headers,
query,
model as string,
psd,
signal as AbortSignal
);
} catch (err) {
return {
response: errResponse(502, err instanceof Error ? err.message : String(err)),
url: this.base(),
headers,
transformedBody: body,
};
}
// Shared per-turn state: plan_item thoughts (cumulative, longest wins).
const order: string[] = [];
const thoughts: Record<string, string> = {};
let sent = 0;
let usage: JsonRecord | null = null;
let errorEvent: JsonRecord | null = null;
const renderNewText = (data: JsonRecord): string => {
const pid = data.id as string | undefined;
if (!pid) return "";
if (!(pid in thoughts)) order.push(pid);
const t = (data.thought as string) || "";
if (t.length >= (thoughts[pid] || "").length) thoughts[pid] = t;
const full = order.map((i) => thoughts[i]).join("");
const piece = full.slice(sent);
sent = full.length;
return piece;
};
if (stream !== false) {
const enc = new TextEncoder();
const sse = new ReadableStream({
start: async (controller) => {
const emit = (obj: JsonRecord) =>
controller.enqueue(enc.encode(`data: ${JSON.stringify(obj)}\n\n`));
let roleEmitted = false;
try {
emit({
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [{ index: 0, delta: { role: "assistant" }, finish_reason: null }],
});
roleEmitted = true;
await this.streamEvents(
headers,
session.sessionId,
session.messageId,
(ev, data) => {
if (ev === "error") {
errorEvent = data;
return true;
}
if (ev === "token_usage") usage = data;
if (ev === "plan_item") {
const piece = renderNewText(data);
if (piece)
emit({
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [{ index: 0, delta: { content: piece }, finish_reason: null }],
});
}
return ev === "done";
},
signal as AbortSignal
);
void roleEmitted;
if (errorEvent) {
emit({
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [],
error: {
message: `trae ${errorEvent.code}: ${errorEvent.message}`,
type: "api_error",
},
});
} else {
emit({
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [{ index: 0, delta: {}, finish_reason: "stop" }],
});
if (usage)
emit({
id: responseId,
object: "chat.completion.chunk",
created,
model,
choices: [],
usage: {
prompt_tokens: usage.prompt_tokens || 0,
completion_tokens: usage.completion_tokens || 0,
total_tokens: usage.total_tokens || 0,
},
});
}
controller.enqueue(enc.encode("data: [DONE]\n\n"));
controller.close();
} catch (err) {
controller.error(err);
}
},
});
return {
response: new Response(sse, {
status: 200,
headers: {
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
Connection: "keep-alive",
},
}),
url: this.base(),
headers,
transformedBody: body,
};
}
// Non-streaming: drive to completion, return chat.completion JSON.
try {
await this.streamEvents(
headers,
session.sessionId,
session.messageId,
(ev, data) => {
if (ev === "error") {
errorEvent = data;
return true;
}
if (ev === "token_usage") usage = data;
if (ev === "plan_item") renderNewText(data);
return ev === "done";
},
signal as AbortSignal
);
} catch (err) {
return {
response: errResponse(502, err instanceof Error ? err.message : String(err)),
url: this.base(),
headers,
transformedBody: body,
};
}
if (errorEvent) {
return {
response: errResponse(502, `trae ${errorEvent.code}: ${errorEvent.message}`),
url: this.base(),
headers,
transformedBody: body,
};
}
const content = order.map((i) => thoughts[i]).join("");
const out: JsonRecord = {
id: responseId,
object: "chat.completion",
created,
model,
choices: [{ index: 0, message: { role: "assistant", content }, finish_reason: "stop" }],
};
if (usage)
out.usage = {
prompt_tokens: usage.prompt_tokens || 0,
completion_tokens: usage.completion_tokens || 0,
total_tokens: usage.total_tokens || 0,
};
return {
response: new Response(JSON.stringify(out), {
status: 200,
headers: { "Content-Type": "application/json" },
}),
url: this.base(),
headers,
transformedBody: body,
};
}
/**
* Headless refresh of the 14-day Cloud-IDE-JWT using the long-lived (~7 month)
* RefreshToken captured during /authorize. Mirrors the desktop client's call to
* POST {apiHost}/cloudide/api/v3/trae/oauth/ExchangeToken
* { ClientID, RefreshToken, ClientSecret: "-", UserID: "" }
* The response uses the same envelope as GetUserToken:
* { ResponseMetadata: { Error?: { Code, Message } }, Result: { Token, RefreshToken,
* TokenExpireAt, RefreshExpireAt, TokenExpireDuration, UserID, TenantID } }
* On Error.Code === "RefreshTokenInvalid" the caller must re-authorize via
* the browser flow — we throw so the connection is marked unusable.
*/
async refreshCredentials(credentials) {
const psd = (credentials?.providerSpecificData as JsonRecord) || {};
const refreshToken = credentials?.refreshToken as string | undefined;
if (!refreshToken) return null;
const host = ((psd.host as string) || "https://api-us-east.trae.ai").replace(/\/$/, "");
const clientId = (psd.clientId as string) || "en1oxy7wnw8j9n";
const url = `${host}/cloudide/api/v3/trae/oauth/ExchangeToken`;
const body = { ClientID: clientId, RefreshToken: refreshToken, ClientSecret: "-", UserID: "" };
const res = await fetch(url, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(body),
});
const text = await res.text();
if (!res.ok) {
throw new Error(`Trae ExchangeToken HTTP ${res.status}: ${text.slice(0, 200)}`);
}
let parsed;
try {
parsed = JSON.parse(text);
} catch {
throw new Error("Trae ExchangeToken: response was not JSON");
}
const errCode = parsed?.ResponseMetadata?.Error?.Code;
if (errCode) {
// Surface invalid-refresh to the caller — BaseExecutor.execute swallows the
// refresh exception, but the next request will hit 401 and trigger fallback;
// we also leave the (stale) accessToken in place so observability shows why.
throw new Error(`Trae ExchangeToken error: ${errCode}`);
}
const result = parsed?.Result;
if (!result?.Token) {
throw new Error("Trae ExchangeToken: response missing Result.Token");
}
return {
accessToken: result.Token as string,
refreshToken: (result.RefreshToken as string) || refreshToken,
expiresAt: result.TokenExpireAt
? new Date(Number(result.TokenExpireAt)).toISOString()
: undefined,
};
}
}
export default TraeExecutor;