diff --git a/eslint.config.js b/eslint.config.js index 66b3f90..8495246 100644 --- a/eslint.config.js +++ b/eslint.config.js @@ -39,6 +39,8 @@ export default [ format: ["camelCase", "PascalCase"], }], + "@typescript-eslint/no-invalid-this": ["error"], + "@stylistic/brace-style": ["error", "allman", { allowSingleLine: true }], "@stylistic/quotes": ["error", "double"], "@stylistic/comma-dangle": ["error", "always-multiline"], diff --git a/package-lock.json b/package-lock.json index d1902fa..f37861c 100644 --- a/package-lock.json +++ b/package-lock.json @@ -58,6 +58,9 @@ }, "engines": { "node": ">=18" + }, + "optionalDependencies": { + "@supabase/supabase-js": "^2.57.2" } }, "node_modules/@alcalzone/ansi-tokenize": { @@ -1552,6 +1555,87 @@ "eslint": ">=9.0.0" } }, + "node_modules/@supabase/auth-js": { + "version": "2.71.1", + "resolved": "https://registry.npmjs.org/@supabase/auth-js/-/auth-js-2.71.1.tgz", + "integrity": "sha512-mMIQHBRc+SKpZFRB2qtupuzulaUhFYupNyxqDj5Jp/LyPvcWvjaJzZzObv6URtL/O6lPxkanASnotGtNpS3H2Q==", + "license": "MIT", + "optional": true, + "dependencies": { + "@supabase/node-fetch": "^2.6.14" + } + }, + "node_modules/@supabase/functions-js": { + "version": "2.4.6", + "resolved": "https://registry.npmjs.org/@supabase/functions-js/-/functions-js-2.4.6.tgz", + "integrity": "sha512-bhjZ7rmxAibjgmzTmQBxJU6ZIBCCJTc3Uwgvdi4FewueUTAGO5hxZT1Sj6tiD+0dSXf9XI87BDdJrg12z8Uaew==", + "license": "MIT", + "optional": true, + "dependencies": { + "@supabase/node-fetch": "^2.6.14" + } + }, + "node_modules/@supabase/node-fetch": { + "version": "2.6.15", + "resolved": "https://registry.npmjs.org/@supabase/node-fetch/-/node-fetch-2.6.15.tgz", + "integrity": "sha512-1ibVeYUacxWYi9i0cf5efil6adJ9WRyZBLivgjs+AUpewx1F3xPi7gLgaASI2SmIQxPoCEjAsLAzKPgMJVgOUQ==", + "license": "MIT", + "optional": true, + "dependencies": { + "whatwg-url": "^5.0.0" + }, + "engines": { + "node": "4.x || >=6.0.0" + } + }, + "node_modules/@supabase/postgrest-js": { + "version": "1.21.4", + "resolved": "https://registry.npmjs.org/@supabase/postgrest-js/-/postgrest-js-1.21.4.tgz", + "integrity": "sha512-TxZCIjxk6/dP9abAi89VQbWWMBbybpGWyvmIzTd79OeravM13OjR/YEYeyUOPcM1C3QyvXkvPZhUfItvmhY1IQ==", + "license": "MIT", + "optional": true, + "dependencies": { + "@supabase/node-fetch": "^2.6.14" + } + }, + "node_modules/@supabase/realtime-js": { + "version": "2.15.5", + "resolved": "https://registry.npmjs.org/@supabase/realtime-js/-/realtime-js-2.15.5.tgz", + "integrity": "sha512-/Rs5Vqu9jejRD8ZeuaWXebdkH+J7V6VySbCZ/zQM93Ta5y3mAmocjioa/nzlB6qvFmyylUgKVS1KpE212t30OA==", + "license": "MIT", + "optional": true, + "dependencies": { + "@supabase/node-fetch": "^2.6.13", + "@types/phoenix": "^1.6.6", + "@types/ws": "^8.18.1", + "ws": "^8.18.2" + } + }, + "node_modules/@supabase/storage-js": { + "version": "2.12.1", + "resolved": "https://registry.npmjs.org/@supabase/storage-js/-/storage-js-2.12.1.tgz", + "integrity": "sha512-QWg3HV6Db2J81VQx0PqLq0JDBn4Q8B1FYn1kYcbla8+d5WDmTdwwMr+EJAxNOSs9W4mhKMv+EYCpCrTFlTj4VQ==", + "license": "MIT", + "optional": true, + "dependencies": { + "@supabase/node-fetch": "^2.6.14" + } + }, + "node_modules/@supabase/supabase-js": { + "version": "2.57.4", + "resolved": "https://registry.npmjs.org/@supabase/supabase-js/-/supabase-js-2.57.4.tgz", + "integrity": "sha512-LcbTzFhHYdwfQ7TRPfol0z04rLEyHabpGYANME6wkQ/kLtKNmI+Vy+WEM8HxeOZAtByUFxoUTTLwhXmrh+CcVw==", + "license": "MIT", + "optional": true, + "dependencies": { + "@supabase/auth-js": "2.71.1", + "@supabase/functions-js": "2.4.6", + "@supabase/node-fetch": "2.6.15", + "@supabase/postgrest-js": "1.21.4", + "@supabase/realtime-js": "2.15.5", + "@supabase/storage-js": "2.12.1" + } + }, "node_modules/@swc-node/core": { "version": "1.14.1", "resolved": "https://registry.npmjs.org/@swc-node/core/-/core-1.14.1.tgz", @@ -1922,6 +2006,13 @@ "form-data": "^4.0.4" } }, + "node_modules/@types/phoenix": { + "version": "1.6.6", + "resolved": "https://registry.npmjs.org/@types/phoenix/-/phoenix-1.6.6.tgz", + "integrity": "sha512-PIzZZlEppgrpoT2QgbnDU+MMzuR6BbCjllj0bM70lWoejMeNJAxCchxnv7J3XFkI8MpygtRpzXrIlmWUBclP5A==", + "license": "MIT", + "optional": true + }, "node_modules/@types/react": { "version": "19.1.12", "resolved": "https://registry.npmjs.org/@types/react/-/react-19.1.12.tgz", @@ -1954,6 +2045,16 @@ "@types/node": "*" } }, + "node_modules/@types/ws": { + "version": "8.18.1", + "resolved": "https://registry.npmjs.org/@types/ws/-/ws-8.18.1.tgz", + "integrity": "sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg==", + "license": "MIT", + "optional": true, + "dependencies": { + "@types/node": "*" + } + }, "node_modules/@typescript-eslint/eslint-plugin": { "version": "8.42.0", "resolved": "https://registry.npmjs.org/@typescript-eslint/eslint-plugin/-/eslint-plugin-8.42.0.tgz", diff --git a/package.json b/package.json index bfa3652..39ad248 100644 --- a/package.json +++ b/package.json @@ -77,6 +77,9 @@ "ulid": "^3.0.1", "write-file-atomic": "^6.0.0" }, + "optionalDependencies": { + "@supabase/supabase-js": "^2.57.2" + }, "devDependencies": { "@sindresorhus/tsconfig": "^7.0.0", "@stylistic/eslint-plugin": "^5.1.0", diff --git a/src/ai/custom-messages.ts b/src/ai/custom-messages.ts index 8ef84c0..24df355 100644 --- a/src/ai/custom-messages.ts +++ b/src/ai/custom-messages.ts @@ -7,7 +7,7 @@ type TCustomMessageType = "error" | "tool-progress"; export type TMessageType = ReturnType | TCustomMessageType; -export class CustomMessage +export abstract class CustomMessage { _type: TCustomMessageType; @@ -21,6 +21,8 @@ export class CustomMessage return this._type; } + abstract isMessageFinished(): boolean; + static fromObject(obj: any): CustomMessage { if (obj._type === "error") @@ -29,7 +31,7 @@ export class CustomMessage } else if (obj._type === "tool-progress") { - return new ToolProgressMessage(obj.toolCall, obj.status, obj.content, obj.confirmState); + return new ToolProgressMessage(obj.toolCall, obj.status, obj.content, obj.confirmState, obj.info); } throw new Error(`Unknown custom message type: ${obj._type}`); @@ -61,6 +63,11 @@ class ErrorMessage extends CustomMessage }; } } + + override isMessageFinished(): boolean + { + return true; + } } class ToolProgressMessage extends CustomMessage @@ -71,12 +78,16 @@ class ToolProgressMessage extends CustomMessage status: "pending" | "pending-confirmation" | "confirmed" | "declined" | "success" | "error"; content?: string; confirmState: typeof ToolProgressMessage.CONFIRM_STATES[number]; + info: { + description?: string; + fileName?: string; + }; static isTypeOf = (m: any) => m instanceof ToolProgressMessage; clone({ status, confirmState }: Pick) { - return new ToolProgressMessage(this.toolCall, status, this.content, confirmState); + return new ToolProgressMessage(this.toolCall, status, this.content, confirmState, this.info); } toggleConfirmState({ direction }: { direction: -1 | 1 }) @@ -124,7 +135,7 @@ class ToolProgressMessage extends CustomMessage return toolCallChunks?.map(i => ToolProgressMessage.createFromChunk(i)).filter(i => !!i) || []; } - constructor(toolCall: ToolCall, status?: ToolProgressMessage["status"], content?: string, confirmState?: ToolProgressMessage["confirmState"]) + constructor(toolCall: ToolCall, status?: ToolProgressMessage["status"], content?: string, confirmState?: ToolProgressMessage["confirmState"], info?: ToolProgressMessage["info"]) { super("tool-progress"); @@ -132,6 +143,12 @@ class ToolProgressMessage extends CustomMessage this.status = status || "pending"; this.content = content; this.confirmState = confirmState || "no"; + this.info = info || {}; + } + + override isMessageFinished(): boolean + { + return this.status !== "pending" && this.status !== "pending-confirmation"; } } @@ -149,9 +166,38 @@ export const isCustomMessage = (message: any): message is CustomMessage => ); }; +const setMessageIsStreaming = (msg: BaseMessage, isStreaming: boolean) => +{ + msg.additional_kwargs["_coba_message_is_streaming"] = isStreaming; + return msg; +}; + +const isMessageStreaming = (msg: TMessage): boolean => +{ + return (msg instanceof BaseMessage) && (msg as BaseMessage).additional_kwargs["_coba_message_is_streaming"] === true; +}; + +const isMessageFinished = (msg: TMessage): boolean => +{ + if (msg instanceof BaseMessage) + { + return !isMessageStreaming(msg); + } + + return msg.isMessageFinished(); +}; + +const resetIsStreamingFlag = (msg: TMessage[]): TMessage[] => +{ + return msg.map(m => (m instanceof BaseMessage) ? setMessageIsStreaming(m, false) : m); +}; + export { ErrorMessage, + isMessageFinished, + isMessageStreaming, + resetIsStreamingFlag, + setMessageIsStreaming, ToolProgressMessage, }; - diff --git a/src/ai/session/chat-server/api-client.ts b/src/ai/session/chat-server/api-client.ts new file mode 100644 index 0000000..876b58b --- /dev/null +++ b/src/ai/session/chat-server/api-client.ts @@ -0,0 +1,357 @@ +// A simple, configurable client library for the documents and links APIs. + +import type { RealtimeChannel, SupabaseClient } from "@supabase/supabase-js"; + +export interface IAuthData +{ + user_id: string; + access_token: string; + api_url: string; + api_key: string; + channel: string; +} + +export interface IClientStatus +{ + type: "consumer-client" | "super-client"; + user_id: string; + external_id: string; + status: "idle" | "working"; +} + +export interface IDocument +{ + id: string; // Internal UUID + external_id: string | null; + user_id: string; + schema_id: string | null; + data: Record; + created_at: string; + updated_at: string; +} + +export interface ILink +{ + id: string; + user_id: string; + from_id: string; + to_id: string; + type: string; + created_at: string; +} + +// Type for creating a new document, with optional external_id and links +export interface INewDocument +{ + data: Record; + external_id?: string; + schema_id?: string | null; + links?: Array<{ to_id: string; type?: string } | { from_id: string; type?: string }>; +} + +export type INewLink = { + from_id: string; + to_id: string; + type?: string; +}; + +export interface IApiClientOptions +{ + baseUrl?: string; + token?: string; +} + +export interface IApiClientCommandListener +{ + handleCommand: (payload: any) => void; +} + +class ApiClient +{ + private baseUrl: string; + private token?: string; + + private authData?: IAuthData; + + private supabase?: SupabaseClient; + private channel?: RealtimeChannel; + + private commandListeners: IApiClientCommandListener[] = []; + + private _initPromise?: Promise; + + constructor(options: IApiClientOptions = {}) + { + this.baseUrl = options.baseUrl || ""; + this.token = options.token; + } + + private async _request(path: string, options: RequestInit): Promise + { + const url = new URL(path, this.baseUrl).href; + const headers = new Headers(options.headers || {}); + headers.set("Content-Type", "application/json"); + + if (this.token) + { + headers.set("Authorization", `Bearer ${this.token}`); + } + + const response = await fetch(url, { ...options, headers }); + + if (!response.ok) + { + const errorBody = await response.json().catch(() => ({ error: `Request failed with status ${response.status}` })); + throw new Error(errorBody.error || "An unknown error occurred"); + } + + const text = await response.text(); + return text ? JSON.parse(text) as T : {} as T; + } + + // --- ---- + addCommandListener(listener: IApiClientCommandListener) + { + if (!this.commandListeners.includes(listener)) + { + this.commandListeners.push(listener); + } + }; + + removeCommandListener(f: IApiClientCommandListener) + { + const i = this.commandListeners.indexOf(f); + + if (i !== -1) + { + this.commandListeners.splice(i, 1); + } + }; + + private pushCommand(payload: any) + { + this.commandListeners.forEach(listener => listener.handleCommand(payload)); + } + + // --- Auth API --- + + getAuthData(): Promise + { + return this._request("api/auth/access-token", { method: "POST" }); + } + + // --- Realtime API --- + + private async createSupabaseClient() + { + const authData = await this.getAuthData(); + + const { createClient } = await import("@supabase/supabase-js"); + + const supabase = createClient(authData.api_url, authData.api_key, { + accessToken: async () => + { + return authData.access_token; + }, + }); + + this.supabase = supabase; + this.authData = authData; + + this.channel = await supabase.channel(authData.channel) + .on( + "postgres_changes", + { + event: "*", + schema: "public", + table: "documents", + filter: `user_id=eq.${authData.user_id}`, + }, + (payload: any) => this.pushCommand(payload), + ) + .on("broadcast", { event: "command" }, (args: any) => this.pushCommand({ args })) + .on("presence", { event: "sync" }, () => this.pushCommand({ presenceState: this.channel?.presenceState() })) + .on("presence", { event: "join" }, (args: any) => this.pushCommand({ args, presenceState: this.channel?.presenceState() })) + .on("presence", { event: "leave" }, (args: any) => this.pushCommand({ args, presenceState: this.channel?.presenceState() })); + + this.channel.subscribe(); + + return this.supabase; + } + + private async ensureSupabaseClient() + { + if (this.supabase) + { + return; + } + + if (!this._initPromise) + { + this._initPromise = this.createSupabaseClient() + .then((_supabase) => + { + // supabase should be ready now! + }) + .catch(error => + { + console.error("Failed to create supabase client", error); + this._initPromise = undefined; + }); + } + + await this._initPromise; + } + + async start() + { + await this.ensureSupabaseClient(); + } + + async setStatus(status: IClientStatus["status"], external_id: IClientStatus["external_id"]) + { + await this.ensureSupabaseClient(); + + await this.channel?.track({ + user_id: this.authData!.user_id, + type: "super-client", + external_id, + status, + } satisfies IClientStatus); + } + + // --- Documents API --- + + listDocuments(): Promise + { + return this._request("api/documents", { method: "GET" }); + } + + /** + * List documents that directly link to the given parent/target document. + * i.e., rows in `links` where links.to_id = targetDocumentId. + * Optionally filter by link `type`. + * + * Note: The server route `GET /api/documents` must support the + * `linked_to` (and optional `type`) query parameters. + */ + listDocumentsLinkingTo(targetDocumentId: string, opts?: { type?: string }): Promise + { + const searchParams = new URLSearchParams({ linked_to: targetDocumentId }); + if (opts?.type) { searchParams.set("type", opts.type); } + return this._request(`api/documents?${searchParams.toString()}`, { method: "GET" }); + } + + /** Convenience alias for readability when treating the target as a parent node. */ + listChildrenOf(parentDocumentId: string, opts?: { type?: string }): Promise + { + return this.listDocumentsLinkingTo(parentDocumentId, opts); + } + + // --- Document access by internal UUID --- + + getDocumentById(id: string): Promise + { + return this._request(`api/documents?id=${id}`, { method: "GET" }); + } + + updateDocumentById(id: string, updates: { data?: Record, schema_id?: string | null }): Promise + { + return this._request("api/documents", { + method: "PUT", + body: JSON.stringify({ id, ...updates }), + }); + } + + deleteDocumentById(id: string): Promise + { + return this._request(`api/documents?id=${id}`, { method: "DELETE" }); + } + + // --- Document access by external_id --- + + getDocumentByExternalId(externalId: string): Promise + { + return this._request(`api/documents/external/${externalId}`, { method: "GET" }); + } + + /** + * Creates a document if it doesn't exist, or updates it if it does. + * @param externalId The external ID of the document to upsert. + * @param documentData The data for the document. + */ + upsertDocument(externalId: string, documentData: { data: Record, schema_id?: string | null }): Promise + { + return this._request(`api/documents/external/${externalId}`, { + method: "PUT", + body: JSON.stringify(documentData), + }); + } + + deleteDocumentByExternalId(externalId: string): Promise + { + return this._request(`api/documents/external/${externalId}`, { method: "DELETE" }); + } + + // --- Document Creation --- + + createDocument(doc: INewDocument): Promise + { + return this._request("api/documents", { + method: "POST", + body: JSON.stringify(doc), + }); + } + + // --- Links API --- + + createLinks(links: INewLink[]): Promise + { + return this._request("api/links", { + method: "POST", + body: JSON.stringify(links), + }); + } + + deleteLinkById(id: string): Promise + { + return this._request(`api/links?id=${id}`, { method: "DELETE" }); + } + + deleteLinkByFromTo(params: { from_id: string; to_id: string; type?: string }): Promise + { + const searchParams = new URLSearchParams({ from_id: params.from_id, to_id: params.to_id }); + if (params.type) + { + searchParams.set("type", params.type); + } + return this._request(`api/links?${searchParams.toString()}`, { method: "DELETE" }); + } + + // --- Broadcast API --- + + broadcast(event: string, payload: { [key: string]: any }) + { + return this.channel?.send({ + type: "broadcast", + event, + payload, + }); + } + + // --- Direct Supabase API --- + + async directUpsertDocument(externalId: string, documentData: { data: Record, schema_id?: string | null }) + { + const rpcParams = { + p_external_id: externalId, + p_data: documentData.data, + p_schema_id: documentData.schema_id, + } as const; + + await this.ensureSupabaseClient(); + await this.supabase!.rpc("upsert_document", rpcParams).select().single(); + } +} + +export { ApiClient }; diff --git a/src/ai/session/chat-server/chat-server-client.ts b/src/ai/session/chat-server/chat-server-client.ts new file mode 100644 index 0000000..e0e0446 --- /dev/null +++ b/src/ai/session/chat-server/chat-server-client.ts @@ -0,0 +1,166 @@ +import { AIMessage, BaseMessage, HumanMessage } from "@langchain/core/messages"; + +import { ErrorMessage, isMessageStreaming, TMessage, TMessageType, ToolProgressMessage } from "../../custom-messages.js"; +import { ChatSession } from "../session.js"; +import { ApiClient } from "./api-client.js"; +import { IChat, IChatMessage } from "./types.js"; + +const URL = process.env["CODE_BANDIT_SERVER_URL"]; +const PAT = process.env["CODE_BANDIT_SERVER_PAT"]; + +const chatServerClient = URL && PAT ? new ApiClient({ baseUrl: URL, token: PAT }) : undefined; + +chatServerClient?.start(); + +function mapMessage(message: TMessage): IChatMessage +{ + const type: TMessageType = message.getType(); + + const base = { + history: [], + historyIndex: 0, + createdAt: Date.now(), + }; + + if (type === "tool-progress") + { + const toolProgressMessage = message as ToolProgressMessage; + + return { + ...base, + role: "tool-progress", + content: toolProgressMessage.content ?? "", + toolCall: { + ...toolProgressMessage.toolCall, + status: toolProgressMessage.status, + confirmState: toolProgressMessage.confirmState, + description: toolProgressMessage.info.description, + fileName: toolProgressMessage.info.fileName, + }, + }; + } + + if (type === "error") + { + const errorMessage = message as ErrorMessage; + + return { + ...base, + role: type, + content: errorMessage.content, + privateData: errorMessage, + }; + } + + if (type === "human") + { + const userMessage = message as HumanMessage; + + return { + ...base, + role: "user", + content: userMessage.text, + }; + } + + if (type === "ai") + { + const assistantMessage = message as AIMessage; + + return { + ...base, + hidden: assistantMessage.text.length === 0, + role: "assistant", + content: assistantMessage.text, + }; + } + + if (type === "system") + { + return { + ...base, + role: type, + content: (message as BaseMessage).text, + }; + } + + return { + ...base, + role: "private", + hidden: true, + content: (message as BaseMessage).text, + privateData: message, + }; +} + +function mapSessionToChat(session: ChatSession): IChat +{ + const chat: IChat = { + id: session.id, + model: { + id: "langchain", + account: { + id: "langchain", + name: "langchain", + }, + name: session.chatServiceOptions.model, + provider: "langchain", + state: { ready: true }, + }, + workspace: { + workDir: session.workDir, + }, + toolMode: session.toolMode, + messages: session.messages.map(mapMessage), + state: session.isWorking ? "working" : "done", + }; + + return chat; +} + +async function pushSessionOnline(session: ChatSession) +{ + await chatServerClient?.directUpsertDocument(session.id, { data: mapSessionToChat(session) }); +} + +async function pushStreamingMessageOnline(session: ChatSession, content: string) +{ + await chatServerClient?.directUpsertDocument(session.id + "/rt", { + data: { + type: "chunk", + external_id: session.id, + chat_id: session.id, + message_number: session.messages.length, + content, + }, + }); +} + +let lastStreamingSync = Date.now(); + +async function syncSession(session: ChatSession) +{ + const aiMessage = session.messages.find(m => isMessageStreaming(m)) as (BaseMessage | undefined); + + if (aiMessage) + { + if (Date.now() > lastStreamingSync + 100) + { + lastStreamingSync = Date.now(); + + await pushStreamingMessageOnline(session, aiMessage.text); + } + + return; + } + + await pushSessionOnline(session); +} + +export +{ + chatServerClient, + mapMessage, + mapSessionToChat, + syncSession, +}; diff --git a/src/ai/session/chat-server/types.ts b/src/ai/session/chat-server/types.ts new file mode 100644 index 0000000..926ee13 --- /dev/null +++ b/src/ai/session/chat-server/types.ts @@ -0,0 +1,242 @@ +type TChatModelProvider = + | "node-llama-cpp" + | "ollama" + | "openai" + | "googleai" + | "groq" + | "xai" + | "openrouter" + | "langchain-js" + | "vercel-ai-sdk" + | string + ; + +export interface IChatModelFeatures +{ + tools?: boolean; + vision?: boolean; + options?: { + temperature?: boolean; + num_ctx?: boolean; + top_k?: boolean; + top_p?: boolean; + min_p?: boolean; + integratedWebSearch?: boolean; + }; +} + +interface IChatModelState +{ + ready?: boolean; + downloadable?: boolean; + removable?: boolean; + hasLocalModelFile?: boolean; +} + +export interface IChatModelOptions +{ + temperature?: number; + num_ctx?: number; + top_k?: number; + top_p?: number; + min_p?: number; + integratedWebSearch?: boolean; +} + +export interface IChatModelConfig +{ + favorite?: boolean; + hidden?: boolean; + options?: IChatModelOptions; +}; + +export interface IChatModel +{ + id: string; + name: string; + provider: TChatModelProvider; + + account: { id: string; name: string; remote?: boolean; removable?: boolean; }; + + displayName?: string; + description?: string; + + modelUri?: string; + modelFile?: string; + + state: IChatModelState; + + contextLength?: number; + size?: number; + parameterSize?: string; + quantizationLevel?: string; + features?: IChatModelFeatures; + + config?: IChatModelConfig; +} + +export class ChatModel implements IChatModel +{ + id: string; + name: string; + provider: TChatModelProvider; + + account: { id: string; name: string; remote?: boolean; removable?: boolean; }; + + displayName?: string; + description?: string; + + modelUri?: string; + modelFile?: string; + + state: IChatModelState; + + contextLength?: number; + size?: number; + parameterSize?: string; + quantizationLevel?: string; + features?: IChatModelFeatures; + + config?: IChatModelConfig; + + constructor(data: IChatModel) + { + this.id = data.id; + this.name = data.name; + this.provider = data.provider; + this.account = structuredClone(data.account); + this.displayName = data.displayName; + this.description = data.description; + this.modelUri = data.modelUri; + this.modelFile = data.modelFile; + this.state = structuredClone(data.state); + this.contextLength = data.contextLength; + this.parameterSize = data.parameterSize; + this.quantizationLevel = data.quantizationLevel; + this.features = structuredClone(data.features); + this.config = structuredClone(data.config); + } +} + +export interface IChatModelInfo +{ + model: Pick; + options?: IChatModelOptions; +} + +export type TAddChatModel = { provider: TChatModelProvider; modelUri: string; startDownload?: boolean; }; + +export type TRemoveChatModel = { provider: TChatModelProvider; modelUri: string; }; + +export interface IAddAccountOpenAI +{ + name: string; + provider: "openai"; + type: "openai"; + apiKey: string; + baseURL?: string; +} + +export interface IAddAccountOpenRouter extends Omit +{ + type: "openrouter"; +} + +interface IAddAccountGoogleAI +{ + name: string; + provider: "googleai"; + apiKey: string; +} + +export type TAddAccount = IAddAccountOpenAI | IAddAccountOpenRouter | IAddAccountGoogleAI; + +export interface TRemoveAccount +{ + provider: TChatModelProvider; + id: string; +} + +export type TChatMessageRole = + | "system" + | "user" + | "assistant" + | "tool" + | "tool-progress" + | "info" + | "error" + | "private" + ; + +export const CONFIRM_STATES = ["no", "yes", "none", "all"] as const; + +export type TConfirmState = typeof CONFIRM_STATES[number]; + +export interface IChatMessageBase +{ + content: string; + images?: string[]; + + toolCall?: { + name: string; + args: Record; + + description?: string; + fileName?: string; + + id?: string; + + status: "pending" | "pending-confirmation" | "confirmed" | "declined" | "success" | "error"; + confirmState?: TConfirmState; + }; + + readonly createdAt: number; + readonly info?: IChatModelInfo; + + showThinking?: boolean; + showDetails?: boolean; + + privateData?: any; +} + +export interface IChatMessage extends IChatMessageBase +{ + role: TChatMessageRole; + hidden?: boolean; + + state?: TChatState; + + history: IChatMessageBase[]; + historyIndex: number; +} + +export interface IWorkspace +{ + workDir?: string; +} + +export type TChatState = "working" | "done" | "stopped"; + +export type TToolMode = "confirm" | "read-only" | "yolo"; + +export interface IChat +{ + id: string; + sourceId?: string; + model: IChatModel; + title?: string; + generatedSummary?: string; + workspace?: IWorkspace; + messages: IChatMessage[]; + currentPrompt?: { content?: string; images?: string[]; messageIndex?: number; }; + state?: TChatState; + toolMode?: TToolMode; + scrollPos?: number; + favorite?: boolean; + useSystemPrompt?: boolean; + showSettings?: boolean; + updatedAt?: number; + deletedAt?: number | "deleted"; +} + +export type TProviderInfo = { [key: string]: any }; diff --git a/src/ai/session/session.ts b/src/ai/session/session.ts index df6adef..1d5c8fa 100644 --- a/src/ai/session/session.ts +++ b/src/ai/session/session.ts @@ -1,8 +1,12 @@ -import { mapChatMessagesToStoredMessages, mapStoredMessageToChatMessage, StoredMessage } from "@langchain/core/messages"; +import { HumanMessage, mapChatMessagesToStoredMessages, mapStoredMessageToChatMessage, StoredMessage } from "@langchain/core/messages"; import { ulid } from "ulid"; -import { IChatServiceOptions } from "../chat-service.js"; -import { CustomMessage, isCustomMessage, TMessage } from "../custom-messages.js"; +import { ChatService, IChatServiceOptions } from "../chat-service.js"; +import { CustomMessage, ErrorMessage, isCustomMessage, isMessageStreaming, resetIsStreamingFlag, TMessage, ToolProgressMessage } from "../custom-messages.js"; +import { work } from "../work.js"; +import { IApiClientCommandListener } from "./chat-server/api-client.js"; +import { chatServerClient, syncSession } from "./chat-server/chat-server-client.js"; +import { TConfirmState } from "./chat-server/types.js"; export type TToolMode = "confirm" | "read-only" | "yolo"; @@ -43,10 +47,10 @@ export function mapSessionToSessionData(session: IChatSession) id: session.id, workDir: session.workDir, toolMode: session.toolMode || "confirm", + streaming: session.streaming, chatServiceOptions: session.chatServiceOptions, systemPrompt: session.systemPrompt, messages: session.messages.map(mapMessageToObject).filter(m => m), - finished: session.finished, }; } @@ -56,10 +60,10 @@ export function mapSessionDataToSession(data: any): IChatSession id: data.id, workDir: data.workDir, toolMode: data.toolMode || "confirm", + streaming: data.streaming || true, chatServiceOptions: data.chatServiceOptions, systemPrompt: data.systemPrompt, messages: data.messages.map(mapObjectToMessage).filter((m: TMessage | undefined) => m), - finished: data.finished, }; } @@ -69,50 +73,55 @@ export interface IChatSession workDir: string; toolMode: TToolMode; + streaming: boolean; chatServiceOptions: IChatServiceOptions; systemPrompt?: string; // TODO: extend into service or template like "%{DEFAULT}% ..." messages: TMessage[]; - finished: number; } -export interface ICreateChatSession extends Omit +export interface ICreateChatSession extends Omit { id?: IChatSession["id"]; messages?: IChatSession["messages"]; - finished?: IChatSession["finished"]; } -type TUpdateListener = (props: { messages: TMessage[]; finished: number; }) => void; +type TUpdateListener = (props: { messages: TMessage[]; working: boolean; toolMode: TToolMode }) => void; -export class ChatSession implements IChatSession +export class ChatSession implements IChatSession, IApiClientCommandListener { id: string; workDir: string; toolMode: TToolMode; chatServiceOptions: IChatServiceOptions; + streaming: boolean; systemPrompt?: string; messages: TMessage[] = []; - finished: number = 0; + + private chatService?: ChatService; storage?: ISessionStorage; private onUpdateListeners: TUpdateListener[] = []; + private abortController?: AbortController; + private constructor(props: IChatSession, storage: ISessionStorage | undefined) { this.id = props.id; this.workDir = props.workDir; this.toolMode = props.toolMode || "confirm"; + this.streaming = props.streaming || true; this.chatServiceOptions = props.chatServiceOptions; this.systemPrompt = props.systemPrompt; this.messages = props.messages || []; - this.finished = Math.min(props.finished || 0, this.messages.length); this.storage = storage; + + chatServerClient?.addCommandListener(this); } static create(props: ICreateChatSession, storage?: ISessionStorage) @@ -121,25 +130,64 @@ export class ChatSession implements IChatSession ...props, id: props.id || ulid(), messages: props.messages || [], - finished: props.finished || 0, }, storage); return chatSession; } - setStorage(storage: ISessionStorage): ChatSession + // TODO: this should not be in session! + handleCommand = (payload: any) => + { + if (payload.new) + { + console.log("PAYLOAD-COMMAND", payload.new.external_id, payload.eventType, JSON.stringify(payload.new.data).length); + } + + if (payload.new?.external_id === `${this.id}/cmd`) + { + if (!this.#isWorking && payload.new.data?.content?.length) + { + this.generateResponse([ + ...this.messages, + new HumanMessage(payload.new.data.content), + ]); + } + } + else if (payload.new?.external_id === `${this.id}/confirm`) + { + if (!this.#isWorking && payload.new.data?.message_index) + { + this.confirmToolUse(payload.new.data.message_index, payload.new.data.confirm_state); + } + } + else if (payload.new?.external_id === `${this.id}/abort`) + { + if (this.#isWorking) + { + this.abort(payload.new.data?.reason || "User cancelled."); + } + } + }; + + setChatService = (chatService: ChatService): ChatSession => + { + this.chatService = chatService; + return this; + }; + + setStorage = (storage: ISessionStorage): ChatSession => { this.storage = storage; return this; - } + }; - private notifyListeners(): void + notifyListeners = (): void => { - const props = { messages: [...this.messages], finished: this.finished }; + const props = { messages: [...this.messages], working: this.#isWorking, toolMode: this.toolMode }; this.onUpdateListeners.forEach(listener => listener(props)); - } + }; - public onUpdate(listener: TUpdateListener): () => void + onUpdate = (listener: TUpdateListener): (() => void) => { this.onUpdateListeners.push(listener); @@ -147,26 +195,141 @@ export class ChatSession implements IChatSession { this.onUpdateListeners = this.onUpdateListeners.filter(l => l !== listener); }; + }; + + #isStreaming = false; + + get isStreaming() + { + return this.#isStreaming; } - async setMessages(messages: TMessage[], finished: number): Promise + setMessages = async (messages: TMessage[]): Promise => { - this.messages = messages; - this.finished = Math.min(finished, this.messages.length); + this.messages = [...messages]; + this.#isStreaming = !!messages.find(m => isMessageStreaming(m)); + + this._save(); + this._saveOnline(); this.notifyListeners(); + }; + + #isWorking = false; + + private set isWorking(isWorking: boolean) + { + this.#isWorking = isWorking; } - private saveQueue: Promise = Promise.resolve(); + get isWorking(): boolean + { + return this.#isWorking; + } + + get isAborted(): boolean + { + return !!this.abortController?.signal?.aborted; + } + + generateResponse = async (messages: TMessage[]) => + { + if (this.#isWorking || !this.chatService) + { + return; + } + + this.#isWorking = true; + this.setMessages(messages); + + this.abortController = new AbortController(); + + work({ + session: this, + chatService: this.chatService, // TODO: can't work use session.chatService? + signal: this.abortController.signal, + }) + .then(messages => + { + this.#isWorking = false; + this.setMessages(messages); + }) + .catch(error => + { + this.#isWorking = false; + const newMessages = [ + ...resetIsStreamingFlag(this.messages), + new ErrorMessage(`ERROR: running work({...}) failed with: ${error.message || error.toString()}`, error), + ]; + this.setMessages(newMessages); + }) + .finally(() => + { + this.abortController = undefined; + }); + }; + + abort = (reason: string) => + { + if (this.abortController && !this.abortController.signal.aborted) + { + this.abortController?.abort(reason); + this.notifyListeners(); + } + }; - async save(): Promise + toggleConfirmState = async (messageIndex: number, direction: -1 | 1): Promise => { - if (!this.storage) + if (ToolProgressMessage.isTypeOf(this.messages[messageIndex])) { - throw new Error("No storage configured!"); + const newMessages = [ + ...this.messages.slice(0, messageIndex), + (this.messages[messageIndex] as ToolProgressMessage).toggleConfirmState({ direction }), + ...this.messages.slice(messageIndex + 1), + ]; + + await this.setMessages(newMessages); + } + }; + + confirmToolUse = async (messageIndex: number, confirmState: TConfirmState): Promise => + { + if (confirmState === "none") + { + this.toolMode = "read-only"; } + else if (confirmState === "all") + { + this.toolMode = "yolo"; + } + + // TODO: reset all other pending tools when toolMode changed! - this.saveQueue = this.saveQueue + this.generateResponse([ + ...this.messages.slice(0, messageIndex), + (this.messages[messageIndex] as ToolProgressMessage).clone({ + status: (confirmState === "yes" || confirmState === "all") ? "confirmed" : "declined", + confirmState, + }), + ...this.messages.slice(messageIndex + 1), + ]); + }; + + setToolMode = (toolMode: TToolMode) => + { + this.toolMode = toolMode; + this.notifyListeners(); + }; + + private _saveQueue: Promise = Promise.resolve(); + private _save = async (): Promise => + { + if (!this.storage || !this.#isStreaming) + { + return Promise.resolve(); + } + + this._saveQueue = this._saveQueue .then(async () => { await this.storage!.saveSession(this); @@ -176,6 +339,30 @@ export class ChatSession implements IChatSession console.error("ERROR: Session save failed!", error); }); - return this.saveQueue; - } + return this._saveQueue; + }; + + private _saveOnlineQueue: Promise = Promise.resolve(); + _saveOnline = (): Promise => + { + this._saveOnlineQueue = this._saveOnlineQueue + .then(async () => + { + await syncSession(this); + }) + .catch((error) => + { + console.error("ERROR: Session save failed!", error); + }); + + return this._saveOnlineQueue; + }; + + flush = async (): Promise => + { + this._saveOnline(); + this._save(); + await this._saveOnlineQueue; + await this._saveQueue; + }; } diff --git a/src/ai/tools/file-system-tools.ts b/src/ai/tools/file-system-tools.ts index 4b978e5..34603de 100644 --- a/src/ai/tools/file-system-tools.ts +++ b/src/ai/tools/file-system-tools.ts @@ -5,7 +5,7 @@ import { globbySync } from "globby"; import path from "path"; import { z } from "zod"; -import { createTool, resolveWithinWorkDir } from "./utils.js"; +import { createTool, resolveWithinWorkDir, setToolProgressMessageInfo } from "./utils.js"; function listDirectory({ directory = "." }: { directory: string }, config?: RunnableConfig): string { @@ -14,6 +14,8 @@ function listDirectory({ directory = "." }: { directory: string }, config?: Runn const resolvedPath = resolveWithinWorkDir(directory, config?.metadata?.["workDir"]); const combinedPath = path.join(resolvedPath, "*"); + setToolProgressMessageInfo({ config, fileName: directory }); + const files = globbySync(combinedPath, { dot: true, onlyFiles: false, @@ -41,6 +43,8 @@ function readFile({ fileName, maxLength }: { fileName: string; maxLength?: numbe { const resolvedPath = resolveWithinWorkDir(fileName, config?.metadata?.["workDir"]); + setToolProgressMessageInfo({ config, fileName }); + const content = readFileSync(resolvedPath).toString().slice(0, maxLength); return content || `The file "${fileName}" is empty.`; @@ -56,7 +60,9 @@ function writeFile({ fileName, fileData }: { fileName: string; fileData: string try { const resolvedPath = resolveWithinWorkDir(fileName, config?.metadata?.["workDir"]); - + + setToolProgressMessageInfo({ config, fileName }); + mkdirSync(path.dirname(resolvedPath), { recursive: true }); writeFileSync(resolvedPath, fileData); @@ -76,6 +82,8 @@ function moveFile({ sourceFileName, destinationFileName }: { sourceFileName: str const resolvedSourcePath = resolveWithinWorkDir(sourceFileName, config?.metadata?.["workDir"]); const resolvedDestinationPath = resolveWithinWorkDir(destinationFileName, config?.metadata?.["workDir"]); + setToolProgressMessageInfo({ config, description: `Rename ${sourceFileName} to ${destinationFileName}` }); + renameSync(resolvedSourcePath, resolvedDestinationPath); return `${sourceFileName} moved to ${destinationFileName}.`; @@ -92,6 +100,8 @@ function deleteFile({ fileName }: { fileName: string }, config: RunnableConfig): { const resolvedPath = resolveWithinWorkDir(fileName, config?.metadata?.["workDir"]); + setToolProgressMessageInfo({ config, fileName }); + unlinkSync(resolvedPath); return `${fileName} deleted.`; @@ -108,6 +118,8 @@ function createDirectory({ fileName }: { fileName: string }, config: RunnableCon { const resolvedPath = resolveWithinWorkDir(fileName, config?.metadata?.["workDir"]); + setToolProgressMessageInfo({ config, fileName }); + mkdirSync(resolvedPath, { recursive: true }); return `${fileName} created.`; @@ -125,6 +137,8 @@ function findFiles({ pattern, directory = "." }: { pattern: string; directory?: const resolvedPath = resolveWithinWorkDir(directory, config?.metadata?.["workDir"]); const combinedPattern = path.join(resolvedPath, pattern); + setToolProgressMessageInfo({ config, fileName: path.join(directory, pattern) }); + const files = globbySync(combinedPattern, { dot: true, onlyFiles: true, @@ -156,6 +170,8 @@ function searchInFiles({ pattern, glob: globPattern, directory = ".", isCaseSens const resolvedPath = resolveWithinWorkDir(directory, workDir); const combinedGlob = path.join(resolvedPath, globPattern); + setToolProgressMessageInfo({ config, description: `Search for ${pattern} in ${path.join(directory, globPattern)}` }); + const files = globbySync(combinedGlob, { dot: true, onlyFiles: true, diff --git a/src/ai/tools/utils.ts b/src/ai/tools/utils.ts index 6fe8e91..236a617 100644 --- a/src/ai/tools/utils.ts +++ b/src/ai/tools/utils.ts @@ -3,6 +3,8 @@ import { tool } from "@langchain/core/tools"; import { realpathSync } from "fs"; import path from "path"; +import { ToolProgressMessage } from "../custom-messages.js"; + export function resolveWithinWorkDir(userPath: string, workDir?: unknown): string { if (!workDir || typeof workDir !== "string") @@ -45,3 +47,14 @@ export function formatEmptyToolOutput(f: Function) { return `Tool \`${funcName(f)}\` produced no output.`; } + +export function setToolProgressMessageInfo(props: { config?: RunnableConfig; description?: string; fileName?: string; }) +{ + const toolProgressMessage = props?.config?.metadata?.["toolProgressMessage"] as (ToolProgressMessage | undefined); + + if (toolProgressMessage) + { + toolProgressMessage.info.description = props.description; + toolProgressMessage.info.fileName = props.fileName; + } +} diff --git a/src/ai/work.ts b/src/ai/work.ts index 955b8af..70bc90d 100644 --- a/src/ai/work.ts +++ b/src/ai/work.ts @@ -3,21 +3,19 @@ import { concat } from "@langchain/core/utils/stream"; import tryCatch from "../utils/try-catch.js"; import { ChatService } from "./chat-service.js"; -import { ErrorMessage, TMessage, ToolProgressMessage } from "./custom-messages.js"; -import { IChatSession } from "./session/session.js"; +import { ErrorMessage, setMessageIsStreaming, TMessage, ToolProgressMessage } from "./custom-messages.js"; +import { ChatSession } from "./session/session.js"; interface WorkProps { chatService: ChatService; - session: IChatSession; - send?: (messages: TMessage[]) => void; - streaming?: boolean; + session: ChatSession; signal?: AbortSignal; } async function work(props: WorkProps) { - const { session, chatService, send, streaming = true, signal } = props; + const { session, chatService, signal } = props; const messages = await workTools({ ...props, session }); @@ -29,7 +27,7 @@ async function work(props: WorkProps) let aiMessage: AIMessageChunk | undefined = undefined; let toolProgressMessages: ToolProgressMessage[] = []; - if (streaming) + if (session.streaming) { const { result: stream, error } = await tryCatch(chatService.stream(session, signal)); @@ -42,10 +40,11 @@ async function work(props: WorkProps) for await (const chunk of stream) { aiMessage = aiMessage !== undefined ? concat(aiMessage, chunk) : chunk; + setMessageIsStreaming(aiMessage, true); if (!aiMessage?.tool_calls?.length && !aiMessage?.tool_call_chunks?.length) { - send?.([...messages, aiMessage]); // TODO: check for race conditions + session.setMessages([...messages, aiMessage]); } else { @@ -53,7 +52,7 @@ async function work(props: WorkProps) if (toolProgressMessages.length) { - send?.([...messages, aiMessage, ...toolProgressMessages]); + session.setMessages([...messages, aiMessage, ...toolProgressMessages]); } } } @@ -73,6 +72,7 @@ async function work(props: WorkProps) if (aiMessage) { + setMessageIsStreaming(aiMessage, false); messages.push(aiMessage); } @@ -83,14 +83,18 @@ async function work(props: WorkProps) // add pending toolCall(s) toolProgressMessages = aiMessage.tool_calls.map(toolCall => new ToolProgressMessage(toolCall)) || []; - messages.push(...toolProgressMessages); + session.setMessages([...messages, ...toolProgressMessages]); - return work({ ...props, session: { ...session, messages } }); + await session.flush(); + + return work(props); } -async function workTools(props: Pick) +async function workTools(props: Pick) { - const { session: { workDir, messages, toolMode }, chatService, send } = props; + const { session, chatService } = props; + + const messages = [...session.messages]; const toolProgressMessages = messages.filter(m => ToolProgressMessage.isTypeOf(m) && (m.status === "pending" || m.status === "confirmed" || m.status === "declined")) as ToolProgressMessage[]; @@ -99,8 +103,6 @@ async function workTools(props: Pick(selectedTool.invoke(toolCall, { metadata })); + const { result, error } = await tryCatch(selectedTool.invoke(toolCall, { + metadata: { + workDir: session.workDir, + toolProgressMessage, + }, + })); if (result) { @@ -165,12 +172,12 @@ async function workTools(props: Pick ToolProgressMessage.isTypeOf(m) && m.status === "pending-confirmation"); } diff --git a/src/commands/chat.tsx b/src/commands/chat.tsx index 78c045f..c7c357c 100644 --- a/src/commands/chat.tsx +++ b/src/commands/chat.tsx @@ -40,15 +40,19 @@ async function initChatSession(options: any) const toolMode: TToolMode = options.readOnly ? "read-only" : options.writeMode ? "yolo" : "confirm"; + const streaming = options.stream !== false; + const session = options.continueSession ? await FileSessionStorage.createFromFile(options.continueSession) - : FileSessionStorage.create({ workDir, toolMode, chatServiceOptions }); + : FileSessionStorage.create({ workDir, toolMode, streaming, chatServiceOptions }); const chatService = new ChatService({ promptLoader: await PromptLoader.create(session), toolProvider: new NodeToolProvider(), }); + session.setChatService(chatService); + return { chatService, session, @@ -63,23 +67,23 @@ async function chat(options: any) chatService, session, startMessage: options.startMessage, - streaming: options.stream !== false, debug: options.debug, }; + session.flush(); + render(, { exitOnCtrlC: false }); } async function exec(message: string, options: any) { - const { chatService, session } = await initChatSession(options); + const { chatService, session } = await initChatSession({ ...options, stream: false }); session.messages.push(new HumanMessage(message)); const messages = await work({ chatService, session, - streaming: false, }); const pendingConfirmation = messages.filter(m => ToolProgressMessage.isTypeOf(m)).find(m => m.status === "pending-confirmation"); @@ -177,3 +181,4 @@ export { addChatCommands, }; + diff --git a/src/ui/chat-app.tsx b/src/ui/chat-app.tsx index 03e1ec0..7672620 100644 --- a/src/ui/chat-app.tsx +++ b/src/ui/chat-app.tsx @@ -5,6 +5,7 @@ import { homedir } from "os"; import React, { useEffect, useState } from "react"; import { ChatService } from "../ai/chat-service.js"; +import { isMessageFinished } from "../ai/custom-messages.js"; import { ChatSession } from "../ai/session/session.js"; import useTerminalSize from "../utils/use-terminal-size.js"; import { useChatController } from "./chat-controller.js"; @@ -19,37 +20,33 @@ interface ChatAppProps chatService: ChatService; session: ChatSession; startMessage?: string; - streaming?: boolean; debug?: boolean; } function ChatApp(props: ChatAppProps) { - const { chatService, session, startMessage, streaming = true, debug } = props; + const { session, startMessage, debug } = props; const { chatServiceOptions } = session; const [_message, setMessage] = useState(""); const { + messages, working, handleInput, action, selected, - chatHistory, - handleSendHistory, tokenUsage, currentGitBranch, } = useChatController({ - chatService, session, - streaming, }); useEffect(() => { if (startMessage) { - handleSendHistory([...chatHistory.messages, new HumanMessage(startMessage)]); + session.generateResponse([...messages, new HumanMessage(startMessage)]); } // eslint-disable-next-line react-hooks/exhaustive-deps }, []); @@ -58,30 +55,32 @@ function ChatApp(props: ChatAppProps) { if (_message.trim() && !working) { - const messages = [...chatHistory.messages, new HumanMessage(_message)]; - handleSendHistory(messages, messages.length); + session.generateResponse([...messages, new HumanMessage(_message)]); setMessage(""); } }; const terminalSize = useTerminalSize(); + let finished = messages.findIndex(m => !isMessageFinished(m)); + if (finished === -1) { finished = messages.length; } + return ( {/* Messages Area */} - + {(message, index) => ( )} - {chatHistory.messages.slice(chatHistory.finished).map((workingItem, index) => ( + {messages.slice(finished).map((workingItem, index) => ( ))} diff --git a/src/ui/chat-controller.ts b/src/ui/chat-controller.ts index 3374364..8dcfc37 100644 --- a/src/ui/chat-controller.ts +++ b/src/ui/chat-controller.ts @@ -2,91 +2,39 @@ import clipboard from "clipboardy"; import { Key, useApp } from "ink"; import { useCallback, useEffect, useState } from "react"; -import { ChatService } from "../ai/chat-service.js"; -import { ErrorMessage, TMessage, ToolProgressMessage } from "../ai/custom-messages.js"; -import { ChatSession } from "../ai/session/session.js"; +import { TMessage, ToolProgressMessage } from "../ai/custom-messages.js"; +import { chatServerClient } from "../ai/session/chat-server/chat-server-client.js"; +import { ChatSession, TToolMode } from "../ai/session/session.js"; import { countTokens, lastMessageFromAI, TTokenUsage } from "../ai/tokens.js"; import { getGitBranch } from "../ai/tools/git-tools.js"; -import { needsToolConfirmation, work } from "../ai/work.js"; interface UseChatControllerProps { - chatService: ChatService; session: ChatSession; - streaming?: boolean; } export function useChatController(props: UseChatControllerProps) { - const { chatService, session, streaming = true } = props; + const { session } = props; const { exit } = useApp(); - const [working, setWorking] = useState(false); - - const [abortController, setAbortController] = useState(); const [ctrlC, setCtrlC] = useState(false); - const [_, setToolMode] = useState(session.toolMode); - - const [chatHistory, setChatHistory] = useState<{ messages: TMessage[]; finished: number }>({ + const [sessionState, setSessionState] = useState<{ messages: TMessage[]; working: boolean; toolMode: TToolMode; }>({ messages: session.messages, - finished: session.finished || 0, + working: session.isWorking, + toolMode: session.toolMode, }); const [tokenUsage, setTokenUsage] = useState(); const [currentGitBranch, setCurrentGitBranch] = useState(); - const selectedIndex = chatHistory.messages.findIndex(m => ToolProgressMessage.isTypeOf(m) && m.status === "pending-confirmation"); + const selectedIndex = sessionState.messages.findIndex(m => ToolProgressMessage.isTypeOf(m) && m.status === "pending-confirmation"); const confirm = selectedIndex !== -1; // This effect subscribes to the session, making it the single source of truth. - useEffect(() => session.onUpdate(setChatHistory), [session]); - - const handleSendHistory = useCallback((messages: TMessage[], finished?: number) => - { - session.setMessages(messages, Math.max(finished || 0, session.finished)); - setWorking(true); - - const abortController = new AbortController(); - setAbortController(abortController); - - work({ - chatService, - session, - streaming, - send: (messages: TMessage[]) => session.setMessages(messages, session.finished), - signal: abortController.signal, - }) - .then(messages => - { - if (needsToolConfirmation(messages)) - { - // Update the session, which will trigger the UI update. - session.setMessages(messages, session.finished); - } - else - { - session.setMessages(messages, messages.length); - } - }) - .catch(error => - { - const newMessages = [ - ...session.messages, - new ErrorMessage(`ERROR: running work({...}) failed with: ${error.message || error.toString()}`, error), - ]; - session.setMessages(newMessages, newMessages.length); - }) - .finally(() => - { - setWorking(false); - }); - }, [ - chatService, - session, - streaming, - ]); + useEffect(() => session.onUpdate(setSessionState), [session]); const handleInput = useCallback((input: string, key: Key): boolean => { @@ -96,10 +44,9 @@ export function useChatController(props: UseChatControllerProps) { setCtrlC(true); } - else if (abortController) + else if (session.isWorking) { - abortController.abort("Ctrl-C"); - setAbortController(undefined); + session.abort("Ctrl-C"); } else { @@ -120,38 +67,13 @@ export function useChatController(props: UseChatControllerProps) { if (key.leftArrow || key.rightArrow) { - const direction = key.leftArrow ? -1 : 1; - - // Instead of setChatHistory, we now update the session directly. - const newMessages = [ - ...session.messages.slice(0, selectedIndex), - (session.messages[selectedIndex] as ToolProgressMessage).toggleConfirmState({ direction }), - ...session.messages.slice(selectedIndex + 1), - ]; - session.setMessages(newMessages, session.finished); + session.toggleConfirmState(selectedIndex, key.leftArrow ? -1 : 1); } - - if (key.return) + else if (key.return) { const confirmState = (session.messages[selectedIndex] as ToolProgressMessage).confirmState; - if (confirmState === "none") - { - setToolMode?.(session.toolMode = "read-only"); - } - else if (confirmState === "all") - { - setToolMode?.(session.toolMode = "yolo"); - } - - handleSendHistory([ - ...session.messages.slice(0, selectedIndex), - (session.messages[selectedIndex] as ToolProgressMessage).clone({ - status: (confirmState === "yes" || confirmState === "all") ? "confirmed" : "declined", - confirmState, - }), - ...session.messages.slice(selectedIndex + 1), - ]); + session.confirmToolUse(selectedIndex, confirmState); } return true; @@ -166,52 +88,44 @@ export function useChatController(props: UseChatControllerProps) if (key.ctrl && input === "w") { - setToolMode(toolMode => session.toolMode = (toolMode === "confirm" ? "read-only" : toolMode === "read-only" ? "yolo" : "confirm")); + session.setToolMode(session.toolMode === "confirm" ? "read-only" : session.toolMode === "read-only" ? "yolo" : "confirm"); return true; } return !!confirm; // do not allow input as long as tool-confirmation is needed }, [ - abortController, confirm, ctrlC, exit, - handleSendHistory, selectedIndex, session, ]); useEffect(() => { - if (!working) + if (!sessionState.working) { - setAbortController(undefined); setCtrlC(false); - } - }, [ - working, - ]); - - useEffect(() => - { - if (!working) - { - session.save(); - setTokenUsage(countTokens(chatHistory.messages)); + setTokenUsage(countTokens(session.messages)); getGitBranch().then(branch => setCurrentGitBranch(branch)); + + chatServerClient?.setStatus("idle", session.id); + } + else + { + chatServerClient?.setStatus("working", session.id); } }, [ - working, - chatHistory.messages, session, + sessionState.working, ]); // TODO: refactor const action = ctrlC - ? working - ? abortController === null + ? sessionState.working + ? session.isAborted ? "Cancelled. Please wait ..." : "Cancel? Press Ctrl-C again ..." : "Quit? Press Ctrl-C again ..." @@ -220,13 +134,11 @@ export function useChatController(props: UseChatControllerProps) : undefined; return { - working, + messages: sessionState.messages, + working: sessionState.working, handleInput, action, selected: selectedIndex, - setToolMode, - chatHistory, - handleSendHistory, tokenUsage, currentGitBranch, };