482 lines
17 KiB
TypeScript
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;
|