@fastino-ai/pioneer-cli 0.2.7 → 0.2.8

package/src/client/WebSocketClient.ts

@@ -42,7 +42,7 @@ export interface WebSocketClientCallbacks {
   onToolCall?: (call: ToolCallRequest) => Promise<string>;
   onAssistantMessage?: (content: string) => void;
   onError?: (error: Error) => void;
-  onDone?: (messages?: HistoryMessage[]) => void;
+  onDone?: (messages?: HistoryMessage[], sessionId?: string) => void;
 }
 
 export interface ChatOptions {
@@ -98,9 +98,10 @@ export class WebSocketClient {
 
       this.ws.on("open", () => {
         this.state = "connected";
-        // For API Gateway WebSocket, send auth message after connecting
-        // (headers aren't forwarded through the request template)
-        if (this.wsUrl.includes("execute-api") && this.apiKey) {
+        // Send auth explicitly as a defensive fallback when header auth is not
+        // propagated by proxy layers (some environments only support
+        // first-message auth for WS).
+        if (this.apiKey) {
           this.ws?.send(JSON.stringify({ type: "auth", api_key: this.apiKey }));
         }
         resolve();
@@ -122,7 +123,7 @@ export class WebSocketClient {
         // Only call onDone if not already called (prevents duplicate)
         if (!this.doneCalled) {
           this.doneCalled = true;
-          this.callbacks.onDone?.();
+          this.callbacks.onDone?.(undefined, undefined);
         }
         if (this.doneResolver) {
           this.doneResolver();
@@ -143,6 +144,15 @@ export class WebSocketClient {
 
   private async handleMessage(message: WebSocketMessage): Promise<void> {
     switch (message.type) {
+      case "auth_required":
+        if (this.apiKey) {
+          this.send({ type: "auth", api_key: this.apiKey });
+        }
+        break;
+
+      case "auth_success":
+        break;
+
       case "stream":
         this.callbacks.onStream?.(message.content as string);
         break;
@@ -211,8 +221,16 @@ export class WebSocketClient {
           this.doneCalled = true;
           // Pass the message history if provided by the backend
           const messages = message.messages as HistoryMessage[] | undefined;
-          console.log("[DEBUG] Done event received, messages:", messages?.length, "roles:", messages?.map(m => m.role));
-          this.callbacks.onDone?.(messages);
+          const sessionId = typeof message.session_id === "string" ? message.session_id : undefined;
+          console.log(
+            "[DEBUG] Done event received, messages:",
+            messages?.length,
+            "roles:",
+            messages?.map((m) => m.role),
+            "sessionId:",
+            sessionId
+          );
+          this.callbacks.onDone?.(messages, sessionId);
         }
         if (this.doneResolver) {
           this.doneResolver();

package/src/config.ts

@@ -36,7 +36,8 @@ export const DEFAULT_BASE_URL = process.env.PIONEER_API_URL || "https://api.pion
 // WebSocket URL environments
 const WS_URLS: Record<string, string> = {
   local: "ws://localhost:5001/mle-agent/ws",
-  dev: "wss://tatd8l0j83.execute-api.us-west-2.amazonaws.com/dev",
+  dev: "wss://api-dev.pioneer.ai/mle-agent/ws",
+  stg: "wss://api-stg.pioneer.ai/mle-agent/ws",
   prod: "wss://api.pioneer.ai/mle-agent/ws",
 };
 
@@ -65,6 +66,12 @@ export function getWsUrl(): string {
     // Dev API Gateway detected - use dev WebSocket URL
     return WS_URLS.dev;
   }
+  if (baseUrl.includes("api-dev.pioneer.ai")) {
+    return WS_URLS.dev;
+  }
+  if (baseUrl.includes("api-stg.pioneer.ai")) {
+    return WS_URLS.stg;
+  }
   if (baseUrl === "https://api.pioneer.ai" || baseUrl === "https://api.pioneer.ai/") {
     // Production detected
     return WS_URLS.prod;