agentlens/packages/opencode-plugin/src/state.ts

import {
  TraceBuilder,
  SpanType,
  SpanStatus,
  DecisionType,
  nowISO,
  getClient,
} from "agentlens-sdk";
import type { JsonValue, TraceStatus } from "agentlens-sdk";
import { extractToolMetadata, safeJsonValue } from "./utils.js";

interface ToolCallState {
  startTime: number;
  tool: string;
  args: unknown;
  sessionID: string;
}

export class SessionState {
  private traces = new Map<string, TraceBuilder>();
  private toolCalls = new Map<string, ToolCallState>();
  private rootSpans = new Map<string, string>();

  startSession(
    sessionId: string,
    metadata?: Record<string, unknown>,
  ): TraceBuilder {
    const trace = new TraceBuilder("opencode-session", {
      sessionId,
      tags: ["opencode", "coding-agent"],
      metadata: metadata ? (safeJsonValue(metadata) as JsonValue) : undefined,
    });

    const rootSpanId = trace.addSpan({
      name: "session",
      type: SpanType.AGENT,
      startedAt: nowISO(),
      metadata: metadata ? (safeJsonValue(metadata) as JsonValue) : undefined,
    });

    this.traces.set(sessionId, trace);
    this.rootSpans.set(sessionId, rootSpanId);
    return trace;
  }

  getTrace(sessionId: string): TraceBuilder | undefined {
    return this.traces.get(sessionId);
  }

  endSession(sessionId: string, status?: TraceStatus): void {
    const trace = this.traces.get(sessionId);
    if (!trace) return;

    const rootSpanId = this.rootSpans.get(sessionId);
    if (rootSpanId) {
      trace.addSpan({
        id: rootSpanId,
        name: "session",
        type: SpanType.AGENT,
        status: status === "ERROR" ? SpanStatus.ERROR : SpanStatus.COMPLETED,
        endedAt: nowISO(),
      });
    }

    trace.end({ status: status ?? "COMPLETED" });
    this.traces.delete(sessionId);
    this.rootSpans.delete(sessionId);
  }

  startToolCall(
    callID: string,
    tool: string,
    args: unknown,
    sessionID: string,
  ): void {
    this.toolCalls.set(callID, {
      startTime: Date.now(),
      tool,
      args,
      sessionID,
    });
  }

  endToolCall(
    callID: string,
    output: string,
    title: string,
    metadata: unknown,
  ): void {
    const call = this.toolCalls.get(callID);
    if (!call) return;
    this.toolCalls.delete(callID);

    const trace = this.traces.get(call.sessionID);
    if (!trace) return;

    const durationMs = Date.now() - call.startTime;
    const rootSpanId = this.rootSpans.get(call.sessionID);
    const toolMeta = extractToolMetadata(call.tool, call.args);

    trace.addSpan({
      name: title || call.tool || "unknown-tool",
      type: SpanType.TOOL_CALL,
      parentSpanId: rootSpanId,
      input: safeJsonValue(call.args),
      output: output as JsonValue,
      durationMs,
      status: SpanStatus.COMPLETED,
      startedAt: new Date(call.startTime).toISOString(),
      endedAt: nowISO(),
      metadata: safeJsonValue({ ...toolMeta, rawMetadata: metadata }),
    });

    const reasoningText =
      title !== call.tool && title
        ? `Selected ${call.tool}: ${title}`
        : `Selected tool: ${call.tool}`;

    trace.addDecision({
      type: DecisionType.TOOL_SELECTION,
      chosen: call.tool as JsonValue,
      alternatives: [],
      reasoning: reasoningText,
      durationMs,
      parentSpanId: rootSpanId,
    });
  }

  recordLLMCall(
    sessionId: string,
    options: {
      model?: { providerID: string; modelID: string };
      agent?: string;
      messageID?: string;
    },
  ): void {
    const trace = this.traces.get(sessionId);
    if (!trace) return;

    const rootSpanId = this.rootSpans.get(sessionId);
    const agentName = options.agent ?? "assistant";
    const modelName = options.model?.modelID ?? "unknown";

    trace.addSpan({
      name: `${agentName} → ${modelName}`,
      type: SpanType.LLM_CALL,
      parentSpanId: rootSpanId,
      status: SpanStatus.COMPLETED,
      startedAt: nowISO(),
      endedAt: nowISO(),
      metadata: safeJsonValue({
        provider: options.model?.providerID,
        model: options.model?.modelID,
        agent: options.agent,
        messageID: options.messageID,
      }),
    });
  }

  recordPermission(
    sessionId: string,
    permission: unknown,
    status: string,
  ): void {
    const trace = this.traces.get(sessionId);
    if (!trace) return;

    const rootSpanId = this.rootSpans.get(sessionId);
    const p = permission as Record<string, unknown> | null;
    const title = (p?.["title"] as string) ?? "permission";
    const permType = (p?.["type"] as string) ?? "unknown";

    trace.addDecision({
      type: DecisionType.ESCALATION,
      chosen: safeJsonValue({ action: status }),
      alternatives: [
        "allow" as JsonValue,
        "deny" as JsonValue,
        "ask" as JsonValue,
      ],
      reasoning: `${permType}: ${title}`,
      parentSpanId: rootSpanId,
    });
  }

  getRootSpanId(sessionId: string): string | undefined {
    return this.rootSpans.get(sessionId);
  }

  getActiveSessionIds(): string[] {
    return Array.from(this.traces.keys());
  }

  flushSession(sessionId: string): void {
    const trace = this.traces.get(sessionId);
    if (!trace) return;

    const rootSpanId = this.rootSpans.get(sessionId);
    if (rootSpanId) {
      trace.addSpan({
        id: rootSpanId,
        name: "session",
        type: SpanType.AGENT,
        status: SpanStatus.COMPLETED,
        endedAt: nowISO(),
      });
    }

    trace.end({ status: "COMPLETED" });

    const transport = getClient();
    if (transport) {
      transport.add(trace.toPayload());
    }

    this.traces.delete(sessionId);
    this.rootSpans.delete(sessionId);
  }
}