pi-langfuse 1.4.0 → 1.4.2

package/README.md

@@ -17,7 +17,7 @@ Langfuse provides open-source observability for LLM applications. This extension
 - **REST fallback for self-hosted Langfuse**: Uses the Langfuse OpenTelemetry SDK first, then verifies that the trace is visible. If a self-hosted OTel ingestion pipeline accepts spans but does not materialize traces, the extension writes the run through Langfuse's REST ingestion API.
 - **Per-Request Generations**: Records a separate `generation` observation for every provider request, including the actual provider payload instead of only the original prompt.
 - **Final Message Capture**: Uses finalized assistant messages for generation and root outputs, so Langfuse shows what the user actually saw in Pi.
-- **Tool Observability**: Creates Langfuse `tool` observations for every tool call, including arguments, results, and error states.
+- **Tool Observability**: Creates Langfuse `tool` observations for every tool call, including arguments, results, error states, and payload/latency metrics.
 - **Parallel Tool Safety**: Correlates tool observations by `toolCallId`, avoiding result mix-ups when Pi runs tools concurrently.
 - **Session Correlation**: Groups traces from the same Pi session under a shared Langfuse session ID.
 - **Cost and Token Tracking**: Records usage and cost details on each generation when Pi/provider payloads expose them.
@@ -32,6 +32,7 @@ Langfuse provides open-source observability for LLM applications. This extension
 - The first generation in a tool-using run can show the assistant's tool-call message, the tool observation shows execution I/O, and the follow-up generation shows the final natural-language answer.
 - Tool failures are marked on the tool observation and reflected in trace-level scores, while later generations still preserve the tool error result in their input history.
 - Shutdown and interrupted runs flush pending telemetry and mark unfinished observations as cancelled/warning instead of silently losing the trace.
+- Agent-end runtime shutdown is deferred so Langfuse flushing does not block Pi's visible turn completion.
 
 ## Prerequisites
 
@@ -265,6 +266,9 @@ Trace (name: "pi-agent")
 | `output` | Tool result, shaped and truncated for readability |
 | `metadata.toolCallId` | Stable Pi tool call identifier |
 | `metadata.isError` | Whether the tool failed |
+| `metadata.durationMs` | Approximate tool runtime in milliseconds |
+| `metadata.inputBytes` | UTF-8 byte size of the shaped tool input payload |
+| `metadata.outputBytes` | UTF-8 byte size of the shaped tool output payload |
 | `level` | `ERROR` for failed tool calls, otherwise `DEFAULT` |
 
 ### Observation-Level Scores

package/index.ts

@@ -10,7 +10,7 @@
 import { basename } from "node:path";
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 
-import { state, resetRunState } from "./src/state.js";
+import { state, resetRunState, runWithSession, setCurrentSession } from "./src/state.js";
 import { ensureConfig, promptForConfig, loadConfig } from "./src/config.js";
 import { shutdownRuntime } from "./src/langfuse.js";
 import { getMessageFromEvent, extractAssistantOutput } from "./src/utils.js";
@@ -51,72 +51,79 @@ export default async function (pi: ExtensionAPI) {
     },
   });
 
-  pi.on("session_start", async (_event, ctx) => {
+  const getSessionId = (ctx?: any) => {
+    try {
+      const sessionFile = ctx?.sessionManager?.getSessionFile?.();
+      return sessionFile ? basename(sessionFile, ".jsonl") : undefined;
+    } catch {
+      return undefined;
+    }
+  };
+
+  const withSession = <T>(ctx: any, fn: () => T): T => runWithSession(getSessionId(ctx) ?? state.currentSessionId, fn);
+
+  pi.on("session_start", async (_event, ctx) => withSession(ctx, async () => {
     state.setupAttemptedThisSession = false;
     await ensureConfig(ctx);
-    const sessionFile = ctx.sessionManager.getSessionFile();
-    if (sessionFile) {
-      state.currentSessionId = basename(sessionFile, ".jsonl");
-    }
     resetRunState();
-  });
+  }));
 
-  pi.on("model_select", async (event) => {
+  pi.on("model_select", async (event, ctx) => withSession(ctx, async () => {
     state.currentModel = event.model?.id || "";
     state.currentProvider = event.model?.provider || "";
-  });
+  }));
 
-  pi.on("before_agent_start", async (event, ctx) => {
+  pi.on("before_agent_start", async (event, ctx) => withSession(ctx, async () => {
     await startAgentRun(event, ctx);
-  });
+  }));
 
-  pi.on("agent_start", async (event, ctx) => {
+  pi.on("agent_start", async (event, ctx) => withSession(ctx, async () => {
     if (!state.agentState?.root) {
       await startAgentRun(event, ctx);
     }
-  });
+  }));
 
-  pi.on("turn_start", async (event) => {
+  pi.on("turn_start", async (event, ctx) => withSession(ctx, async () => {
     await startTurnObservation(event);
-  });
+  }));
 
-  pi.on("before_provider_request", async (event) => {
+  pi.on("before_provider_request", async (event, ctx) => withSession(ctx, async () => {
     await startGeneration(event);
-  });
+  }));
 
-  pi.on("after_provider_response", async (event) => {
+  pi.on("after_provider_response", async (event, ctx) => withSession(ctx, async () => {
     updateGenerationMetadata(event);
-  });
+  }));
 
-  pi.on("message_update", async (event) => {
+  pi.on("message_update", async (event, ctx) => withSession(ctx, async () => {
     recordTTFT(event);
     const message = getMessageFromEvent(event);
     if (message?.role === "assistant" && state.agentState) {
       state.agentState.latestAssistantOutput = extractAssistantOutput(message);
     }
-  });
+  }));
 
-  pi.on("message_end", async (event) => {
+  pi.on("message_end", async (event, ctx) => withSession(ctx, async () => {
     await finishGenerationFromMessage(event);
-  });
+  }));
 
-  pi.on("tool_execution_start", async (event) => {
+  pi.on("tool_execution_start", async (event, ctx) => withSession(ctx, async () => {
     await startToolObservation(event);
-  });
+  }));
 
-  pi.on("tool_call", async (event) => {
+  pi.on("tool_call", async (event, ctx) => withSession(ctx, async () => {
     await startToolObservation(event);
-  });
+  }));
 
-  pi.on("tool_result", async (event) => {
+  pi.on("tool_result", async (event, ctx) => withSession(ctx, async () => {
     await finishToolObservation(event);
-  });
+  }));
 
-  pi.on("tool_execution_end", async (event) => {
+  pi.on("tool_execution_end", async (event, ctx) => withSession(ctx, async () => {
     await finishToolObservation(event);
-  });
+  }));
 
-  pi.on("turn_end", async (event) => {
+  pi.on("turn_end", async (event, ctx) => withSession(ctx, async () => {
     state.turnCount++;
     const message = getMessageFromEvent(event);
     if (message?.role === "assistant") {
@@ -124,12 +131,16 @@ export default async function (pi: ExtensionAPI) {
       await finishGenerationFromMessage(event);
     }
     finishTurnObservation(event);
-  });
+  }));
 
-  pi.on("agent_end", async (event) => {
+  pi.on("agent_end", async (event, ctx) => withSession(ctx, async () => {
     await finishAgentRun(event);
-    await shutdownRuntime();
-  });
+    setTimeout(() => {
+      shutdownRuntime().catch((error) => {
+        console.warn("📊 Langfuse: Deferred shutdown failed", error);
+      });
+    }, 0);
+  }));
 
   const handleSessionInterruption = (reason: string) => {
     if (state.agentState?.root) {
@@ -139,15 +150,21 @@ export default async function (pi: ExtensionAPI) {
     resetRunState();
   };
 
-  pi.on("session_before_switch", async () => {
-    handleSessionInterruption("Session switched");
+  pi.on("session_before_switch", async (_event, ctx) => {
+    const sessionId = getSessionId(ctx);
+    if (sessionId) {
+      setCurrentSession(sessionId);
+    }
   });
 
-  pi.on("session_before_fork", async () => {
-    handleSessionInterruption("Session forked");
+  pi.on("session_before_fork", async (_event, ctx) => {
+    const sessionId = getSessionId(ctx);
+    if (sessionId) {
+      setCurrentSession(sessionId);
+    }
   });
 
-  pi.on("session_compact", async (event) => {
+  pi.on("session_compact", async (event, ctx) => withSession(ctx, async () => {
     if (state.agentState?.root) {
       const parent = state.agentState.activeTurn ?? state.agentState.root;
       try {
@@ -165,10 +182,10 @@ export default async function (pi: ExtensionAPI) {
         // ignore
       }
     }
-  });
+  }));
 
-  pi.on("session_shutdown", async () => {
+  pi.on("session_shutdown", async (_event, ctx) => withSession(ctx, async () => {
     handleSessionInterruption("Session shutdown before agent completed");
     await shutdownRuntime();
-  });
+  }));
 }

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "pi-langfuse",
-	"version": "1.4.0",
+	"version": "1.4.2",
 	"description": "Langfuse extension for Pi coding agent",
 	"repository": {
 		"type": "git",
@@ -12,6 +12,16 @@
 	"homepage": "https://github.com/gooyoung/pi-langfuse#readme",
 	"type": "module",
 	"main": "index.ts",
+	"files": [
+		"index.ts",
+		"src/",
+		"types/",
+		"README.md",
+		"README_CN.md",
+		"image.png",
+		"skills-lock.json",
+		"tsconfig.json"
+	],
 	"scripts": {
 		"typecheck": "tsc --noEmit"
 	},

package/src/handlers/tool.ts

@@ -7,6 +7,7 @@ import {
   shapePayload,
   extractTextContent,
   truncate,
+  estimatePayloadBytes,
 } from "../utils.js";
 import { MAX_TOOL_PAYLOAD_LENGTH } from "../constants.js";
 
@@ -22,27 +23,36 @@ export async function startToolObservation(event: Record<string, unknown>) {
 
   try {
     const toolName = getToolName(event);
+    const toolInput = getToolInput(event);
+    const shapedInput = shapePayload(toolInput, { maxString: MAX_TOOL_PAYLOAD_LENGTH });
+    const inputBytes = estimatePayloadBytes(shapedInput, MAX_TOOL_PAYLOAD_LENGTH);
     const parent = state.agentState.activeTurn ?? state.agentState.root;
     const tool = parent.startObservation
       ? parent.startObservation(
           toolName,
           {
-            input: shapePayload(getToolInput(event), { maxString: MAX_TOOL_PAYLOAD_LENGTH }),
-            metadata: { toolName, toolCallId },
+            input: shapedInput,
+            metadata: { toolName, toolCallId, inputBytes },
           },
           { asType: "tool" },
         )
       : (await getRuntime()).startObservation(
           toolName,
           {
-            input: shapePayload(getToolInput(event), { maxString: MAX_TOOL_PAYLOAD_LENGTH }),
-            metadata: { toolName, toolCallId },
+            input: shapedInput,
+            metadata: { toolName, toolCallId, inputBytes },
           },
           { asType: "tool" },
         );
 
     state.toolCallCount++;
-    state.agentState.activeTools.set(toolCallId, { observation: tool, toolName, ended: false });
+    state.agentState.activeTools.set(toolCallId, {
+      observation: tool,
+      toolName,
+      ended: false,
+      startedAt: Date.now(),
+      inputBytes,
+    });
   } catch (e) {
     console.warn("📊 Langfuse: Failed to start tool observation", e);
   }
@@ -73,15 +83,22 @@ export async function finishToolObservation(event: Record<string, unknown>) {
     event;
 
   try {
+    const shapedOutput = shapePayload(output, { maxString: MAX_TOOL_PAYLOAD_LENGTH });
+    const outputBytes = estimatePayloadBytes(shapedOutput, MAX_TOOL_PAYLOAD_LENGTH);
+    const durationMs = Math.max(0, Date.now() - activeTool.startedAt);
+
     activeTool.observation
       .update({
-        output: shapePayload(output, { maxString: MAX_TOOL_PAYLOAD_LENGTH }),
+        output: shapedOutput,
         level: isError ? "ERROR" : "DEFAULT",
         statusMessage: isError ? truncate(String(event.error ?? output), 1_000) : undefined,
         metadata: {
           toolName: activeTool.toolName,
           toolCallId,
           isError,
+          durationMs,
+          inputBytes: activeTool.inputBytes,
+          outputBytes,
         },
       })
       .end();

package/src/state.ts

@@ -1,38 +1,141 @@
+import { AsyncLocalStorage } from "node:async_hooks";
 import type { Config, AgentState } from "./types.js";
 
+export interface SessionRunState {
+  currentModel: string;
+  currentProvider: string;
+  agentState: AgentState | null;
+  toolCallCount: number;
+  errorCount: number;
+  turnCount: number;
+}
+
+const DEFAULT_SESSION_ID = "__pi_langfuse_default_session__";
+
+let activeSessionId = DEFAULT_SESSION_ID;
+const sessionScope = new AsyncLocalStorage<string>();
+
+function createSessionRunState(): SessionRunState {
+  return {
+    currentModel: "",
+    currentProvider: "",
+    agentState: null,
+    toolCallCount: 0,
+    errorCount: 0,
+    turnCount: 0,
+  };
+}
+
+function normalizeSessionId(sessionId?: string) {
+  return sessionId || DEFAULT_SESSION_ID;
+}
+
+function getActiveSessionId() {
+  return sessionScope.getStore() ?? activeSessionId;
+}
+
+export function getSessionRunState(sessionId = getActiveSessionId()): SessionRunState {
+  const normalizedSessionId = normalizeSessionId(sessionId);
+  let sessionState = state.sessionStates.get(normalizedSessionId);
+  if (!sessionState) {
+    sessionState = createSessionRunState();
+    state.sessionStates.set(normalizedSessionId, sessionState);
+  }
+  return sessionState;
+}
+
+export function setCurrentSession(sessionId?: string) {
+  activeSessionId = normalizeSessionId(sessionId);
+  getSessionRunState(activeSessionId);
+}
+
+export function runWithSession<T>(sessionId: string | undefined, fn: () => T): T {
+  const normalizedSessionId = normalizeSessionId(sessionId);
+  setCurrentSession(normalizedSessionId);
+  return sessionScope.run(normalizedSessionId, fn);
+}
+
 export const state = {
   config: null as Config | null,
   setupAttemptedThisSession: false,
-  
-  currentSessionId: "",
-  currentModel: "",
-  currentProvider: "",
-  agentState: null as AgentState | null,
-  
-  // Evaluation tracking state
-  toolCallCount: 0,
-  errorCount: 0,
-  turnCount: 0,
+  sessionStates: new Map<string, SessionRunState>(),
+
+  get currentSessionId() {
+    const sessionId = getActiveSessionId();
+    return sessionId === DEFAULT_SESSION_ID ? "" : sessionId;
+  },
+  set currentSessionId(sessionId: string) {
+    setCurrentSession(sessionId);
+  },
+
+  get currentModel() {
+    return getSessionRunState().currentModel;
+  },
+  set currentModel(model: string) {
+    getSessionRunState().currentModel = model;
+  },
+
+  get currentProvider() {
+    return getSessionRunState().currentProvider;
+  },
+  set currentProvider(provider: string) {
+    getSessionRunState().currentProvider = provider;
+  },
+
+  get agentState() {
+    return getSessionRunState().agentState;
+  },
+  set agentState(agentState: AgentState | null) {
+    getSessionRunState().agentState = agentState;
+  },
+
+  get toolCallCount() {
+    return getSessionRunState().toolCallCount;
+  },
+  set toolCallCount(toolCallCount: number) {
+    getSessionRunState().toolCallCount = toolCallCount;
+  },
+
+  get errorCount() {
+    return getSessionRunState().errorCount;
+  },
+  set errorCount(errorCount: number) {
+    getSessionRunState().errorCount = errorCount;
+  },
+
+  get turnCount() {
+    return getSessionRunState().turnCount;
+  },
+  set turnCount(turnCount: number) {
+    getSessionRunState().turnCount = turnCount;
+  },
 };
 
-export function resetRunState() {
-  state.agentState = null;
-  state.toolCallCount = 0;
-  state.errorCount = 0;
-  state.turnCount = 0;
-  state.currentModel = "";
-  state.currentProvider = "";
+export function resetRunState(sessionId = getActiveSessionId()) {
+  state.sessionStates.set(normalizeSessionId(sessionId), createSessionRunState());
 }
 
-export function computeEvaluationScores() {
-  const toolSuccessRate = state.toolCallCount > 0 ? (state.toolCallCount - state.errorCount) / state.toolCallCount : 1;
-  const sessionHadErrors = state.errorCount > 0;
+export function clearAllSessionStates() {
+  state.sessionStates.clear();
+  activeSessionId = DEFAULT_SESSION_ID;
+  getSessionRunState();
+}
+
+export function computeEvaluationScores(sessionId = getActiveSessionId()) {
+  const sessionState = getSessionRunState(sessionId);
+  const toolSuccessRate =
+    sessionState.toolCallCount > 0
+      ? (sessionState.toolCallCount - sessionState.errorCount) / sessionState.toolCallCount
+      : 1;
+  const sessionHadErrors = sessionState.errorCount > 0;
 
   return {
-    tool_call_count: state.toolCallCount,
-    turn_count: state.turnCount,
-    total_tool_errors: state.errorCount,
+    tool_call_count: sessionState.toolCallCount,
+    turn_count: sessionState.turnCount,
+    total_tool_errors: sessionState.errorCount,
     tool_success_rate: toolSuccessRate,
     session_had_errors: sessionHadErrors ? 1 : 0,
   };
 }
+
+getSessionRunState();

package/src/types.ts

@@ -87,6 +87,8 @@ export interface ToolState {
   observation: LangfuseObservation;
   toolName: string;
   ended: boolean;
+  startedAt: number;
+  inputBytes: number;
 }
 
 export interface AgentState {

package/src/utils.ts

@@ -97,6 +97,10 @@ export function safeSerialize(value: unknown, maxLength = MAX_TOOL_PAYLOAD_LENGT
   }
 }
 
+export function estimatePayloadBytes(value: unknown, maxLength = MAX_TOOL_PAYLOAD_LENGTH): number {
+  return new TextEncoder().encode(safeSerialize(value, maxLength)).length;
+}
+
 export function extractTextContent(content: unknown, maxLength?: number): string | undefined {
   if (typeof content === "string") {
     return maxLength ? truncate(content, maxLength) : content;

package/.agents/skills/langfuse/SKILL.md

@@ -1,140 +0,0 @@
----
-name: langfuse
-description: Interact with Langfuse and access its documentation. Use when needing to (1) query or modify Langfuse data programmatically via the CLI — traces, prompts, datasets, scores, sessions, and any other API resource, (2) look up Langfuse documentation, concepts, integration guides, or SDK usage, or (3) understand how any Langfuse feature works. This skill covers CLI-based API access (via npx) and multiple documentation retrieval methods.
-allowed-tools:
-  - WebFetch(domain:langfuse.com)
-  - Bash(curl *langfuse.com/*)
-  - Bash(npx langfuse-cli api __schema *)
-  - Bash(npx langfuse-cli api * --help *)
-  - Bash(npx langfuse-cli api * list *)
-  - Bash(npx langfuse-cli api * get *)
-  - Bash(bunx langfuse-cli api __schema *)
-  - Bash(bunx langfuse-cli api * --help *)
-  - Bash(bunx langfuse-cli api * list *)
-  - Bash(bunx langfuse-cli api * get *)
----
-
-# Langfuse
-
-This skill helps you use Langfuse effectively across all common workflows: instrumenting applications, migrating prompts, debugging traces, and accessing data programmatically.
-
-## Core Principles
-
-Follow these principles for ALL Langfuse work:
-
-1. **Documentation First**: NEVER implement based on memory. Always fetch current docs before writing code (Langfuse updates frequently) See the section below on how to access documentation.
-2. **CLI for Data Access**: Use `langfuse-cli` when querying/modifying Langfuse data. See the section below on how to use the CLI. 
-3. **Best Practices by Use Case**: Check the relevant reference file below for use-case-specific guidelines before implementing
-4. **Use latest Langfuse versions**: Unless the user specified otherwise or there's a good reason, always use the latest version of Langfuse SDKs/APIs.
-
-
-## Use case specific references
-
-- instrumenting an existing function/application: references/instrumentation.md
-- migrating prompts from a codebase into Langfuse: references/prompt-migration.md
-- capturing user feedback (thumbs, ratings, implicit signals) as scores on traces: references/user-feedback.md
-- further tips on using the Langfuse CLI: references/cli.md
-- upgrading or migrating Langfuse SDKs to the latest version: references/sdk-upgrade.md
-- systematic error analysis — reading traces, building failure taxonomy, deciding what to fix: references/error-analysis.md
-- submitting feedback about this skill: references/skill-feedback.md
-
-## 1. Langfuse API via CLI
-
-Use the `langfuse-cli` to interact with the full Langfuse REST API from the command line. Run via npx (no install required):
-
-Start by discovering the schema and available arguments:
-
-```bash
-# Discover all available resources
-npx langfuse-cli api __schema
-
-# List actions for a resource
-npx langfuse-cli api <resource> --help
-
-# Show args/options for a specific action
-npx langfuse-cli api <resource> <action> --help
-```
-
-### Credentials
-
-Set environment variables before making calls:
-
-```bash
-export LANGFUSE_PUBLIC_KEY=pk-lf-...
-export LANGFUSE_SECRET_KEY=sk-lf-...
-export LANGFUSE_HOST=https://cloud.langfuse.com # example for EU cloud. For US cloud it's us.cloud.langfuse.com, and can also be a self-hosted URL. The server must always be specified in order to access Langfuse.
-```
-
-If not set, ask the user to set them in their shell or a `.env` file (do not ask them to paste keys into chat for security reasons). Keys are found in Langfuse UI → Settings → API Keys.
-
-### Detailed CLI Reference
-
-For common workflows, tips, and full usage patterns, see [references/cli.md](references/cli.md).
-
-## 2. Langfuse Documentation
-
-Three methods to access Langfuse docs, in order of preference. **Always prefer your application's native web fetch and search tools** (e.g., `WebFetch`, `WebSearch`, `mcp_fetch`, etc.) over `curl` when available. The URLs and patterns below work with any fetching method — the `curl` examples are just illustrative.
-
-### 2a. Documentation Index (llms.txt)
-
-Fetch the full index of all documentation pages:
-
-```bash
-curl -s https://langfuse.com/llms.txt
-```
-
-Returns a structured list of every doc page with titles and URLs. Use this to discover the right page for a topic, then fetch that page directly.
-
-Alternatively, you can start on `https://langfuse.com/docs` and explore the site to find the page you need.
-
-### 2b. Fetch Individual Pages as Markdown
-
-Any page listed in llms.txt can be fetched as markdown by appending `.md` to its path or by using `Accept: text/markdown` in the request headers. Use this when you know which page contains the information needed. Returns clean markdown with code examples and configuration details.
-
-```bash
-curl -s "https://langfuse.com/docs/observability/overview.md"
-curl -s "https://langfuse.com/docs/observability/overview" -H "Accept: text/markdown"
-```
-
-### 2c. Search Documentation
-
-When you need to find information across all docs and github issues/discussions without knowing the specific page:
-
-```bash
-curl -s "https://langfuse.com/api/search-docs?query=<url-encoded-query>"
-```
-
-Example:
-
-```bash
-curl -s "https://langfuse.com/api/search-docs?query=How+do+I+trace+LangGraph+agents"
-```
-
-Returns a JSON response with:
-
-- `query`: the original query
-- `answer`: a JSON string containing an array of matching documents, each with:
-  - `url`: link to the doc page
-  - `title`: page title
-  - `source.content`: array of relevant text excerpts from the page
-
-Search is a great fallback if you cannot find the relevant pages or need more context. Especially useful when debugging issues as all GitHub Issues and Discussions are also indexed. Responses can be large — extract only the relevant portions.
-
-### Documentation Workflow
-
-1. Start with **llms.txt** to orient — scan for relevant page titles
-2. **Fetch specific pages** when you identify the right one
-3. Fall back to **search** when the topic is unclear and you want more context
-
-## Skill Feedback
-
-When the user expresses that something about this skill is not working as expected, gives incorrect guidance, is missing information, or could be improved — offer to submit feedback to the Langfuse skill maintainers. This includes when:
-
-- The skill gave wrong or outdated instructions
-- A workflow didn't produce the expected result
-- The user wishes the skill covered something it doesn't
-- The user explicitly says something like "this should work differently" or "this is wrong"
-
-**Do NOT trigger this** for issues with Langfuse itself (the product) — only for issues with this skill's instructions and behavior.
-
-When triggered, follow the process in [references/skill-feedback.md](references/skill-feedback.md).