@xalia/agent 0.6.9 → 0.6.10

package/src/chat/server/chatContextManager.ts

@@ -34,15 +34,14 @@ import {
   ISessionFileManager,
   ISessionFileManagerEventHandler,
   createSessionFilesManagerPrompt,
-} from "./sessionFileManager";
-import {
   SessionFileDescriptor,
   SessionFileEntry,
-} from "../data/dbSessionFileModels";
+} from "../../agent/tools/fileManager";
 // eslint-disable-next-line max-len
 import { ContextTransactionWithWorkspace } from "../../agent/contextWithWorkspace";
 import { getErrorString } from "./errorUtils";
 import { createUserMessage } from "../../agent/agent";
+import { TokenCounter } from "../../agent/tokenCounter";
 
 const logger = getLogger();
 
@@ -64,7 +63,8 @@ export interface ICheckpointWriter {
 export class ChatContextTransaction implements IContextTransaction {
   private readonly baseTx: ContextTransactionWithWorkspace;
   private readonly sessionUUID: string;
-  /// Index of final message in the committed context
+  /// Index of final message in the committed context.  If this has changed
+  /// before we try to commit this tx, the commit will fail.
   private readonly baseMsgIdx: number | undefined;
   private readonly startingLLMContextLength: number;
   private readonly pendingMessages: ConversationMessage[];
@@ -216,7 +216,6 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
   private readonly llmContext: CompressingContextManager;
 
   private nextMessageIdx: number;
-  private pendingUserMessages: ServerUserMessage[];
 
   // Compression state
   private readonly checkpointWriter: ICheckpointWriter;
@@ -226,6 +225,10 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
   private readonly fileManager: ISessionFileManager;
   private fileManagerDescriptionsDirty: boolean;
 
+  // LLM and token counting
+  private readonly llm: ILLM;
+  private tokenCounter: TokenCounter;
+
   constructor(
     systemPrompt: string,
     sessionMessages: SessionMessage[],
@@ -254,7 +257,6 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
       defaultUserName,
       sessionUUID
     );
-    this.pendingUserMessages = [];
     this.llmContext = new CompressingContextManager(
       systemPrompt,
       llmMessages,
@@ -266,13 +268,14 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
     this.fileManager = fileManager;
     fileManager.addEventHandler(this);
     this.fileManagerDescriptionsDirty = true;
+    this.llm = llm;
+    this.tokenCounter = new TokenCounter(llm.getModel());
   }
 
   // IContextManager.getLLMContext
   getLLMContext(): MessageParam[] {
     if (this.fileManagerDescriptionsDirty) {
       const prompt = createSessionFilesManagerPrompt(this.fileManager);
-      logger.debug(`[ChatContextManager] filemanager prompt:\n${prompt}`);
       this.llmContext.setPromptFragment("file_manager", prompt);
       this.fileManagerDescriptionsDirty = false;
     }
@@ -324,7 +327,22 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
 
   // Get the conversation (to send to clients)
   getConversationMessages(): ConversationMessage[] {
-    return this.conversationMessages.concat(this.pendingUserMessages);
+    return this.conversationMessages;
+  }
+
+  // Get current context usage (tokens)
+  getContextUsage(): { used: number; max: number } {
+    // Update tokenCounter if model changed
+    const currentModel = this.llm.getModel();
+    if (this.tokenCounter.getModel() !== currentModel) {
+      this.tokenCounter.free();
+      this.tokenCounter = new TokenCounter(currentModel);
+    }
+
+    const messages = this.getLLMContext();
+    const used = this.tokenCounter.countMessagesTokens(messages);
+    const max = this.tokenCounter.getContextWindow();
+    return { used, max };
   }
 
   processUserMessage(
@@ -336,7 +354,11 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
     // pass in our generated messages back into `startAgentResponse`.
 
     // Filter out null messages immediately.
-    if (!msg.imageB64 && !msg.message) {
+    if (
+      !msg.imageB64 &&
+      !msg.message &&
+      (!msg.attachedFiles || msg.attachedFiles.length === 0)
+    ) {
       return undefined;
     }
 
@@ -352,8 +374,13 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
     if (msg.imageB64) {
       userMessage.imageB64 = msg.imageB64;
     }
+    if (msg.attachedFiles) {
+      userMessage.attachedFiles = msg.attachedFiles;
+    }
+    if (msg.race_mode) {
+      userMessage.race_mode = msg.race_mode;
+    }
 
-    this.pendingUserMessages.push(userMessage);
     return userMessage;
   }
 
@@ -373,32 +400,18 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
     // Sanity check the state - the incoming user messages should match the
     // pending user messages.
 
-    const numMessages = this.pendingUserMessages.length;
-    assert(numMessages > 0);
-    if (msgs.length !== this.pendingUserMessages.length) {
-      throw new Error(
-        `length mismatch: msgs: ${JSON.stringify(msgs)}, ` +
-          `this.pendingUserMsgs: ${JSON.stringify(this.pendingUserMessages)}`
-      );
-    }
-    assert(msgs[0].message_idx === this.pendingUserMessages[0].message_idx);
-    assert(
-      msgs[numMessages - 1].message_idx ===
-        this.pendingUserMessages[numMessages - 1].message_idx
-    );
-
+    const numMessages = msgs.length;
+    assert(numMessages > 0, "no messages");
     // Collect the pending user messages and allocate a starting index for
     // agent messages and tool calls.
 
-    const pendingUserMessages = this.pendingUserMessages;
-    const baseMsgIdx = this.lastMessageIdx();
+    const baseMsgIdx = this.lastCommittedMessageIdx();
     const curAgentMsgIdx = this.getNextMessageIdx();
-    this.pendingUserMessages = [];
 
     // Compute the new llm messages
 
     const llmUserMessages: UserMessageParam[] = [];
-    for (const msg of pendingUserMessages) {
+    for (const msg of msgs) {
       const userMsg = createUserMessage(
         msg.message,
         msg.imageB64,
@@ -419,12 +432,19 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
       end: false,
     };
 
+    // Update file manager fragment BEFORE starting the transaction
+    if (this.fileManagerDescriptionsDirty) {
+      const prompt = createSessionFilesManagerPrompt(this.fileManager);
+      this.llmContext.setPromptFragment("file_manager", prompt);
+      this.fileManagerDescriptionsDirty = false;
+    }
+
     const baseTx = await this.llmContext.startTx(llmUserMessages);
     const contextTx = new ChatContextTransaction(
       baseTx,
       this.sessionUUID,
       baseMsgIdx,
-      pendingUserMessages,
+      msgs,
       curAgentMsgIdx
     );
     return { agentFirstChunk, contextTx };
@@ -432,7 +452,7 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
 
   async endAgentResponse(tx: IContextTransaction): Promise<SessionMessage[]> {
     assert(tx instanceof ChatContextTransaction);
-    if (tx.baseMessageIdx() !== this.lastMessageIdx()) {
+    if (tx.baseMessageIdx() !== this.lastCommittedMessageIdx()) {
       throw new Error(
         `Tx stale? tx.baseMessageIdx=${String(tx.baseMessageIdx())}, ` +
           `this.conv: ${JSON.stringify(this.conversationMessages)}`
@@ -449,12 +469,13 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
     const newLLMMessages = tx.newMessages();
 
     const messageListError = (error: string) => {
-      throw new Error(
-        `${error}:` +
-          `\n  newSessionMessages: ${JSON.stringify(newSessionMessages)}` +
-          `\n  pending: ${JSON.stringify(pending)}` +
-          `\n  newLLMMessages: ${JSON.stringify(newLLMMessages)}`
-      );
+      const fullError =
+        `[endAgentResponse] Message list validation failed - ${error}:` +
+        `\n  newSessionMessages: ${JSON.stringify(newSessionMessages)}` +
+        `\n  pending: ${JSON.stringify(pending)}` +
+        `\n  newLLMMessages: ${JSON.stringify(newLLMMessages)}`;
+      logger.error(fullError);
+      throw new Error(fullError);
     };
 
     if (newSessionMessages.length !== numPending) {
@@ -506,7 +527,7 @@ export class ChatContextManager implements ISessionFileManagerEventHandler {
     return newSessionMessages;
   }
 
-  private lastMessageIdx(): number | undefined {
+  private lastCommittedMessageIdx(): number | undefined {
     const numMsgs = this.conversationMessages.length;
     if (numMsgs > 0) {
       return this.conversationMessages[numMsgs - 1].message_idx;

package/src/chat/server/conversation.ts

@@ -127,6 +127,9 @@ export function userMessageToConversationMessage(
   if (userMsgData.imageB64) {
     msg.imageB64 = userMsgData.imageB64;
   }
+  if (userMsgData.attachedFiles) {
+    msg.attachedFiles = userMsgData.attachedFiles;
+  }
   return msg;
 }
 

package/src/chat/server/imageGeneratorTools.ts

@@ -4,15 +4,16 @@ import { AgentEx, IAgentToolProvider, ToolCallResult } from "../../agent/agent";
 import { ToolDescriptor } from "../../agent/llm";
 import { ImageGenerator } from "../../agent/imageGenerator";
 import {
-  ChatSessionFileManager,
+  ISessionFileManager,
   ToolCallResultWithFileRef,
-} from "./sessionFileManager";
-import { makeParseArgsFn } from "./tools";
-import { getSessionFileMimeTypeFromDataUrl } from "../data/dbSessionFileModels";
+  getSessionFileMimeTypeFromDataUrl,
+  makeParseArgsFn,
+} from "../../agent/tools";
 import { getOpenAIClientParams } from "./openAIRouterLLM";
 import { DEFAULT_IMAGE_GEN_MODEL, ImageGenLLM } from "../../agent/imageGenLLM";
 import { createSpecializedLLM } from "../../agent/agentUtils";
 import { IPlatform } from "../../agent/iplatform";
+import { CANVAS_WORKSPACE_FILENAME } from "../constants";
 
 const logger = getLogger();
 
@@ -95,13 +96,17 @@ export async function genImageTool(
 
 export async function genImageFileTool(
   platform: IPlatform,
-  fileManager: ChatSessionFileManager
+  fileManager: ISessionFileManager
 ): Promise<IAgentToolProvider> {
   const GEN_IMAGE_FILE_DESC: ToolDescriptor = {
     type: "function",
     function: {
       name: "gen_image_file",
-      description: "Generate image (into session file manager)",
+      description:
+        "Generate or edit an image. Saves result to file manager. " +
+        `Canvas workspace is '${CANVAS_WORKSPACE_FILENAME}'. ` +
+        "When editing canvas images, set add_to_canvas=true to " +
+        "place the result back onto the canvas.",
       parameters: {
         type: "object",
         properties: {
@@ -119,7 +124,14 @@ export async function genImageFileTool(
           },
           input_name: {
             type: "string",
-            description: "(Optional) input image filename",
+            description:
+              "(Optional) input image filename. " +
+              `Use '${CANVAS_WORKSPACE_FILENAME}' for canvas edits.`,
+          },
+          add_to_canvas: {
+            type: "boolean",
+            description:
+              "Set true to place result on canvas (use when editing canvas)",
           },
         },
         required: ["prompt", "name", "summary"],
@@ -128,26 +140,33 @@ export async function genImageFileTool(
   } as const;
 
   const imageGenerator = await createImageGenerator(platform);
-  const getPromptNameSummary = makeParseArgsFn(
+  const getArgs = makeParseArgsFn(
     ["prompt", "name", "summary"] as const,
-    ["input_name"] as const
+    ["input_name"] as const,
+    ["add_to_canvas"] as const
   );
   const toolFn = async (
     _: AgentEx,
     args: unknown
   ): Promise<ToolCallResultWithFileRef> => {
-    const { prompt, name, summary, input_name } = getPromptNameSummary(args);
+    const { prompt, name, summary, input_name, add_to_canvas } = getArgs(args);
+    const addToCanvas = add_to_canvas ?? false;
+
     const input_image = input_name
       ? await fileManager.getFileContent(input_name)
       : undefined;
 
     const image = await imageGenerator.generate(prompt, input_image);
-    const mimeType = getSessionFileMimeTypeFromDataUrl(image);
     await fileManager.putFileContent(name, summary, image);
     const uri = fileManager.getSessionFileRelativeUrl(name);
+
     return {
       response: uri,
-      _meta: { "xalia/fileUri": uri, "xalia/fileMimeType": mimeType },
+      _meta: {
+        "xalia/fileUri": uri,
+        "xalia/fileMimeType": getSessionFileMimeTypeFromDataUrl(image),
+        ...(addToCanvas && { "xalia/addToCanvas": "true" }),
+      },
     };
   };
 

package/src/chat/server/openAIRouterLLM.ts

@@ -20,7 +20,6 @@ const DEFAULT_PROVIDER_URLS: Record<string, string | undefined> = {
   openrouter: "https://openrouter.ai/api/v1",
   openai: "https://api.openai.com/v1",
   anthropic: "https://api.anthropic.com/v1",
-  together: "https://api.together.xyz/v1",
 };
 
 const DEFAULT_MODEL_MAP: Record<string, string | undefined> = {
@@ -35,8 +34,6 @@ const DEFAULT_MODEL_MAP: Record<string, string | undefined> = {
   "anthropic/claude-sonnet-4": "openrouter",
   "anthropic/claude-sonnet-4.5": "openrouter",
   "claude-3-7-sonnet-20250219": "anthropic",
-  "arcee-ai/AFM-4.5B-Preview": "together",
-  "meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": "together",
 };
 
 type CachedMaps = {
@@ -158,7 +155,7 @@ export class OpenAIRouterLLM implements ILLM {
     tools?: ToolDescriptor[],
     onMessage?: (msg: string, end: boolean) => Promise<void>,
     onReasoning?: (reasoning: string) => Promise<void>
-  ): Promise<Completion> {
+  ): Promise<{ stop: (msg: string) => void; completion: Promise<Completion> }> {
     return OpenAILLMStreaming.makeRequest(
       this.openai,
       this.model,