import { defineWebSocket, useCache, usePlatform } from '@kb-labs/sdk';
import { SessionManager } from '@kb-labs/agent-core';

// src/ws/session-stream-handler.ts
var CACHE_PREFIX = "agent:run:";
var EVENT_TOPIC_PREFIX = "agent:events:";
var CACHE_TTL = 36e5;
var RunManagerImpl = class {
  /** Live agents and listeners (not cacheable) */
  activeRuns = /* @__PURE__ */ new Map();
  /** Session-level listeners: sessionId → Set<callback> — receive events from ALL runs in session */
  sessionListeners = /* @__PURE__ */ new Map();
  /** sessionId per runId — set when run is registered */
  runSessionMap = /* @__PURE__ */ new Map();
  /**
   * Generate unique run ID
   */
  generateRunId() {
    return `run-${Date.now()}-${Math.random().toString(36).slice(2, 7)}`;
  }
  /**
   * Register a new run
   */
  async register(runId, task, agent, sessionManager, sessionId) {
    const now = (/* @__PURE__ */ new Date()).toISOString();
    const run = {
      runId,
      task,
      status: "pending",
      agent,
      sessionManager,
      startedAt: now,
      listeners: /* @__PURE__ */ new Set(),
      lastSeq: 0,
      eventBuffer: []
    };
    this.activeRuns.set(runId, run);
    if (sessionId) {
      this.runSessionMap.set(runId, sessionId);
    }
    await this.saveToCache(run);
    return run;
  }
  /**
   * Get run by ID (from memory first, then cache for state)
   */
  get(runId) {
    return this.activeRuns.get(runId);
  }
  /**
   * Check if run exists (in memory or cache)
   */
  async exists(runId) {
    if (this.activeRuns.has(runId)) {
      return true;
    }
    const state = await this.getState(runId);
    return state !== null;
  }
  /**
   * Get run state from cache (for completed runs or cross-process access)
   */
  async getState(runId) {
    const cache = useCache();
    if (!cache) {
      return null;
    }
    return cache.get(`${CACHE_PREFIX}${runId}`);
  }
  /**
   * Update run status
   */
  async updateStatus(runId, status, extra) {
    const run = this.activeRuns.get(runId);
    if (run) {
      run.status = status;
      if (extra) {
        Object.assign(run, extra);
      }
      await this.saveToCache(run);
    }
  }
  /**
   * Save run state to cache
   */
  async saveToCache(run) {
    const cache = useCache();
    if (!cache) {
      return;
    }
    const state = {
      runId: run.runId,
      task: run.task,
      status: run.status,
      startedAt: run.startedAt,
      completedAt: run.completedAt,
      durationMs: run.durationMs,
      summary: run.summary,
      error: run.error
    };
    await cache.set(`${CACHE_PREFIX}${run.runId}`, state, CACHE_TTL);
  }
  /** Track subscriptions for cleanup */
  subscriptions = /* @__PURE__ */ new Map();
  /**
   * Add event listener to run (uses eventBus for cross-process)
   */
  addListener(runId, callback) {
    const platform = usePlatform();
    if (!platform?.eventBus) {
      const run2 = this.activeRuns.get(runId);
      if (run2) {
        run2.listeners.add(callback);
        return true;
      }
      return false;
    }
    const topic = `${EVENT_TOPIC_PREFIX}${runId}`;
    const unsubscribe = platform.eventBus.subscribe(topic, async (event) => {
      try {
        callback(event);
      } catch {
      }
    });
    if (!this.subscriptions.has(runId)) {
      this.subscriptions.set(runId, /* @__PURE__ */ new Map());
    }
    this.subscriptions.get(runId).set(callback, unsubscribe);
    const run = this.activeRuns.get(runId);
    if (run) {
      run.listeners.add(callback);
    }
    return true;
  }
  /**
   * Remove event listener from run
   */
  removeListener(runId, callback) {
    const runSubs = this.subscriptions.get(runId);
    if (runSubs) {
      const unsubscribe = runSubs.get(callback);
      if (unsubscribe) {
        unsubscribe();
        runSubs.delete(callback);
      }
      if (runSubs.size === 0) {
        this.subscriptions.delete(runId);
      }
    }
    const run = this.activeRuns.get(runId);
    if (run) {
      run.listeners.delete(callback);
    }
  }
  /**
   * Broadcast event to all listeners of a run (uses eventBus for cross-process)
   * Assigns monotonic sequence number for reliable ordering
   * @returns The event with seq assigned (for persistence)
   */
  broadcast(runId, event) {
    const platform = usePlatform();
    const run = this.activeRuns.get(runId);
    let seqEvent = event;
    if (run) {
      run.lastSeq++;
      seqEvent = { ...event, seq: run.lastSeq, runId };
      run.eventBuffer.push(seqEvent);
    }
    if (platform?.eventBus) {
      const topic = `${EVENT_TOPIC_PREFIX}${runId}`;
      platform.eventBus.publish(topic, seqEvent).catch(() => {
      });
    }
    if (run) {
      for (const listener of run.listeners) {
        try {
          listener(seqEvent);
        } catch {
        }
      }
    }
    const sessionId = this.runSessionMap.get(runId);
    if (sessionId) {
      const sListeners = this.sessionListeners.get(sessionId);
      if (sListeners) {
        for (const listener of sListeners) {
          try {
            listener(seqEvent);
          } catch {
          }
        }
      }
    }
    return seqEvent;
  }
  /**
   * Add a session-level listener — receives events from ALL runs in this session.
   */
  addSessionListener(sessionId, callback) {
    if (!this.sessionListeners.has(sessionId)) {
      this.sessionListeners.set(sessionId, /* @__PURE__ */ new Set());
    }
    this.sessionListeners.get(sessionId).add(callback);
  }
  /**
   * Remove a session-level listener.
   */
  removeSessionListener(sessionId, callback) {
    const listeners = this.sessionListeners.get(sessionId);
    if (listeners) {
      listeners.delete(callback);
      if (listeners.size === 0) {
        this.sessionListeners.delete(sessionId);
      }
    }
  }
  /**
   * Get replay buffer for a run (all events emitted since start).
   * Used by WS handler to send missed events on late connection.
   * Optionally filter by afterSeq to only get events the client hasn't seen.
   */
  getEventBuffer(runId, afterSeq) {
    const run = this.activeRuns.get(runId);
    if (!run) {
      return [];
    }
    if (afterSeq != null) {
      return run.eventBuffer.filter((e) => e.seq != null && e.seq > afterSeq);
    }
    return [...run.eventBuffer];
  }
  /**
   * List all active runs
   */
  listActive() {
    return Array.from(this.activeRuns.values()).map((r) => ({
      runId: r.runId,
      task: r.task,
      status: r.status,
      startedAt: r.startedAt
    }));
  }
  /**
   * Request graceful stop of a running agent (and its child agents via propagated AbortSignal).
   * Agent finishes its current tool call then exits at the next iteration boundary.
   */
  requestStop(runId) {
    const run = this.activeRuns.get(runId);
    if (!run || run.status !== "running" || !run.agent) {
      return false;
    }
    run.agent.requestStop();
    return true;
  }
  /**
   * Clean up completed runs from memory (cache handles its own TTL)
   */
  cleanup() {
    for (const [runId, run] of this.activeRuns) {
      if (run.status === "completed" || run.status === "failed" || run.status === "stopped") {
        this.activeRuns.delete(runId);
      }
    }
  }
};
var GLOBAL_KEY = "__kb_agent_run_manager__";
if (!globalThis[GLOBAL_KEY]) {
  globalThis[GLOBAL_KEY] = new RunManagerImpl();
}
var RunManager = globalThis[GLOBAL_KEY];
function getTurnIdFromEvent(event) {
  if ((event.type === "agent:start" || event.type === "agent:end" || event.type === "agent:error") && !event.parentAgentId && event.agentId) {
    return `turn-${event.agentId}`;
  }
  return null;
}
function getTurnSignature(turn) {
  return `${turn.id}:${turn.status}:${turn.completedAt || ""}:${turn.steps.length}`;
}
async function resolveTurnForEvent(sessionManager, sessionId, event) {
  const turns = await sessionManager.getTurns(sessionId);
  if (turns.length === 0) {
    return null;
  }
  const explicitTurnId = getTurnIdFromEvent(event);
  if (explicitTurnId) {
    const explicit = turns.find((turn) => turn.id === explicitTurnId);
    if (explicit) {
      return explicit;
    }
  }
  const assistantTurns = turns.filter((turn) => turn.type === "assistant").sort((a, b) => b.sequence - a.sequence);
  return assistantTurns[0] || null;
}
var connectionState = /* @__PURE__ */ new WeakMap();
var session_stream_handler_default = defineWebSocket({
  path: "/session/:sessionId",
  description: "Persistent session event stream (all runs)",
  handler: {
    async onConnect(ctx, sender) {
      const sessionId = ctx.hostContext.params?.sessionId;
      if (!sessionId) {
        await sender.send({
          type: "error",
          payload: { code: "MISSING_SESSION_ID", message: "Session ID is required" },
          timestamp: Date.now()
        });
        sender.close(4e3, "Missing session ID");
        return;
      }
      ctx.platform.logger.info(`[session-ws] Client connected to session ${sessionId}`);
      const sessionManager = new SessionManager(ctx.cwd);
      const lastTurnSignatures = /* @__PURE__ */ new Map();
      await sender.send({
        type: "connection:ready",
        payload: { runId: sessionId, connectedAt: (/* @__PURE__ */ new Date()).toISOString() },
        timestamp: Date.now()
      });
      try {
        const snapshot = await sessionManager.getConversationSnapshot(sessionId);
        await sender.send({
          type: "conversation:snapshot",
          payload: {
            sessionId,
            completedTurns: snapshot.completedTurns,
            activeTurns: snapshot.activeTurns,
            totalTurns: snapshot.totalTurns,
            timestamp: (/* @__PURE__ */ new Date()).toISOString()
          },
          timestamp: Date.now()
        });
        ctx.platform.logger.info(
          `[session-ws] Sent snapshot: ${snapshot.completedTurns.length} completed + ${snapshot.activeTurns.length} active turns`
        );
      } catch (err) {
        ctx.platform.logger.error(`[session-ws] Failed to send snapshot: ${err}`);
      }
      const eventCallback = async (event) => {
        const evtSessionId = event.sessionId || event.metadata?.sessionId;
        if (evtSessionId && evtSessionId !== sessionId) {
          return;
        }
        const targetSessionId = evtSessionId || sessionId;
        try {
          const turn = await resolveTurnForEvent(sessionManager, targetSessionId, event);
          if (turn) {
            const signature = getTurnSignature(turn);
            const previousSignature = lastTurnSignatures.get(turn.id);
            if (previousSignature !== signature) {
              lastTurnSignatures.set(turn.id, signature);
              await sender.send({
                type: "turn:snapshot",
                payload: { sessionId: targetSessionId, turn, sequenceNumber: turn.sequence },
                timestamp: Date.now()
              });
            }
          }
        } catch (err) {
          ctx.platform.logger.error(`[session-ws] Failed to process turn snapshot: ${err}`);
        }
        if (event.type === "agent:end" && !event.parentAgentId) {
          const runId = event.runId || event.metadata?.runId;
          await sender.send({
            type: "run:completed",
            payload: {
              runId: runId ?? "unknown",
              success: event.data.success,
              summary: event.data.summary,
              durationMs: event.data.durationMs
            },
            timestamp: Date.now()
          });
        }
      };
      connectionState.set(ctx, { sessionId, callback: eventCallback });
      RunManager.addSessionListener(sessionId, eventCallback);
    },
    async onMessage(ctx, message, sender) {
      if (message.type === "ping") {
        const sessionId = connectionState.get(ctx)?.sessionId ?? "";
        await sender.send({
          type: "connection:ready",
          payload: { runId: sessionId ?? "", connectedAt: (/* @__PURE__ */ new Date()).toISOString() },
          timestamp: Date.now()
        });
      }
    },
    async onDisconnect(ctx) {
      const state = connectionState.get(ctx);
      if (state) {
        RunManager.removeSessionListener(state.sessionId, state.callback);
        connectionState.delete(ctx);
      }
      ctx.platform.logger.info(`[session-ws] Client disconnected from session ${state?.sessionId}`);
    },
    async onError(ctx, error, sender) {
      ctx.platform.logger.error(`[session-ws] WebSocket error: ${error.message}`);
      await sender.send({
        type: "error",
        payload: { code: "INTERNAL_ERROR", message: error.message },
        timestamp: Date.now()
      });
    }
  }
});

export { session_stream_handler_default as default };
//# sourceMappingURL=session-stream-handler.js.map
//# sourceMappingURL=session-stream-handler.js.map