From 875ba62fb2030c762b7cd06f6d6a48ea12708106 Mon Sep 17 00:00:00 2001
From: zerafachris <christopher.zerafa@blocklabs.io>
Date: Wed, 15 Jul 2026 11:38:37 +0200
Subject: [PATCH] feat(hook): custom trace tags and session grouping via
 commands

Adds two opt-in, related config options for attributing traces to external
work items, both empty by default (no behavior change for existing users):

- CC_LANGFUSE_TAG_COMMAND: run once per hook run; each non-empty stdout line
  is merged into the trace tags (de-duped, order-preserving).
- CC_LANGFUSE_SESSION_LABEL_COMMAND + CC_LANGFUSE_SESSION_LABEL_MODE: use the
  first stdout line as a label to group the session's traces
  (prefix / collapse / off).

One script can drive both -- tags use every line, the label uses the first.
Session labeling changes only the trace grouping id; the incremental-read
state key stays on the raw session id, so relabeling never re-emits past
turns and needs no state migration. Both paths are fail-open and bounded
(timeout, tag count/length, 199-char session id cap).

Adds tests/unit/test_custom_tags.py (15 cases). Full suite green.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .claude-plugin/plugin.json     |  22 ++++++
 README.md                      |  44 ++++++++++++
 hooks/langfuse_hook.py         | 123 +++++++++++++++++++++++++++++++-
 tests/unit/test_custom_tags.py | 125 +++++++++++++++++++++++++++++++++
 4 files changed, 312 insertions(+), 2 deletions(-)
 create mode 100644 tests/unit/test_custom_tags.py

diff --git a/.claude-plugin/plugin.json b/.claude-plugin/plugin.json
index 35c551c..43d0d6c 100644
--- a/.claude-plugin/plugin.json
+++ b/.claude-plugin/plugin.json
@@ -61,6 +61,28 @@
       "type": "string",
       "title": "Trace ID seed",
       "description": "Optional. Derive deterministic trace ids as Langfuse.create_trace_id(seed=f\"{seed}:{turn_number}\") so external systems can precompute them. Use a unique seed per session (e.g. the --session-id UUID). Leave empty for auto-generated ids."
+    },
+    "CC_LANGFUSE_TAG_COMMAND": {
+      "type": "string",
+      "title": "Custom tag command",
+      "description": "Optional shell command run once per hook run; each non-empty line of its stdout is added to the trace tags. Use it to attribute traces to a git branch, ticket id, CI run, or cost center without editing the hook. Fail-open and bounded (see CC_LANGFUSE_TAG_TIMEOUT; max 20 tags of 64 chars). Example: a script that echoes 'sc-1234' derived from the current branch."
+    },
+    "CC_LANGFUSE_TAG_TIMEOUT": {
+      "type": "number",
+      "title": "Custom tag command timeout (seconds)",
+      "description": "Timeout for CC_LANGFUSE_TAG_COMMAND and CC_LANGFUSE_SESSION_LABEL_COMMAND. Default 2. If a command runs longer it is killed and contributes nothing.",
+      "default": 2
+    },
+    "CC_LANGFUSE_SESSION_LABEL_COMMAND": {
+      "type": "string",
+      "title": "Session label command",
+      "description": "Optional shell command whose first non-empty stdout line becomes a label used to group this session's traces in Langfuse (e.g. a ticket id from the branch). The same script can serve CC_LANGFUSE_TAG_COMMAND (tags use every line, the label uses the first). Fail-open and bounded."
+    },
+    "CC_LANGFUSE_SESSION_LABEL_MODE": {
+      "type": "string",
+      "title": "Session label mode",
+      "description": "How the label groups sessions. 'prefix' (default): session id becomes '<label>/<session>', one Langfuse session per Claude session grouped under the label. 'collapse': every session sharing a label merges into one. 'off': disable. Affects only trace grouping; incremental read state stays keyed on the raw session id.",
+      "default": "prefix"
     }
   }
 }
diff --git a/README.md b/README.md
index 49a6416..7d93596 100644
--- a/README.md
+++ b/README.md
@@ -138,6 +138,50 @@ For how Langfuse Cloud handles data it receives, see the Langfuse privacy policy
 https://langfuse.com/privacy . When using a self-hosted Langfuse instance, your data
 stays within your own infrastructure.
 
+## Custom tags
+
+Set `CC_LANGFUSE_TAG_COMMAND` to a shell command and every trace gets tagged with each
+non-empty line the command prints to stdout. Use it to attribute traces to whatever your
+team tracks — a git branch, a ticket, a CI run, a cost center — without editing the hook.
+
+```bash
+CC_LANGFUSE_TAG_COMMAND='~/bin/langfuse-tags.sh'
+```
+
+```bash
+#!/usr/bin/env bash
+# ~/bin/langfuse-tags.sh — tag traces with the ticket id on the current branch
+branch=$(git rev-parse --abbrev-ref HEAD 2>/dev/null) || exit 0
+[[ $branch =~ [sS][cC]-([0-9]+) ]] && { echo "sc-${BASH_REMATCH[1]}"; echo "layer:build"; }
+```
+
+The command runs once per hook run (its result applies to every turn emitted in that run,
+so within a session the tags stay stable). It's fail-open — a missing command, a non-zero
+exit, a timeout, or a crash simply adds no tags and never interrupts tracing — and bounded:
+`CC_LANGFUSE_TAG_TIMEOUT` seconds (default 2), at most 20 tags of 64 characters each.
+Because the logic lives in your own script, it survives plugin updates.
+
+### Grouping sessions by a label
+
+`CC_LANGFUSE_SESSION_LABEL_COMMAND` takes the first non-empty line of a command's stdout and
+uses it to group a session's traces in Langfuse — so, for example, all Claude Code work on
+one ticket lands under one grouping key.
+
+```bash
+CC_LANGFUSE_SESSION_LABEL_COMMAND='~/bin/langfuse-tags.sh'   # first line = the label
+CC_LANGFUSE_SESSION_LABEL_MODE='prefix'                      # prefix | collapse | off
+```
+
+- `prefix` (default): session id becomes `<label>/<session>` — one Langfuse session per
+  Claude session, grouped under the label.
+- `collapse`: every session sharing a label merges into a single Langfuse session.
+- `off`: disabled.
+
+The same script can drive both features — `CC_LANGFUSE_TAG_COMMAND` uses every line as a tag,
+`CC_LANGFUSE_SESSION_LABEL_COMMAND` uses the first line as the label. Point both at one script
+that prints the ticket id on line 1. Labeling affects only trace grouping; the incremental-read
+state stays keyed on the raw session id, so it never re-emits past turns.
+
 ## Reconfigure
 
 In Claude Code, run:
diff --git a/hooks/langfuse_hook.py b/hooks/langfuse_hook.py
index f4563c9..1efeeef 100644
--- a/hooks/langfuse_hook.py
+++ b/hooks/langfuse_hook.py
@@ -15,6 +15,7 @@
 import os
 import random
 import sys
+import subprocess
 import threading
 import time
 import hashlib
@@ -48,6 +49,17 @@ def _opt(name: str) -> str:
 DEBUG = _opt("CC_LANGFUSE_DEBUG").lower() == "true"
 SKILL_TAGS = (_opt("CC_LANGFUSE_SKILL_TAGS") or "true").lower() == "true"
 CAPTURE_SKILL_CONTENT = _opt("CC_LANGFUSE_CAPTURE_SKILL_CONTENT").lower() == "true"
+CUSTOM_TAG_COMMAND = _opt("CC_LANGFUSE_TAG_COMMAND")
+SESSION_LABEL_COMMAND = _opt("CC_LANGFUSE_SESSION_LABEL_COMMAND")
+SESSION_LABEL_MODE = (_opt("CC_LANGFUSE_SESSION_LABEL_MODE") or "prefix").lower()
+try:
+    CUSTOM_TAG_TIMEOUT = float(_opt("CC_LANGFUSE_TAG_TIMEOUT") or "2")
+except ValueError:
+    CUSTOM_TAG_TIMEOUT = 2.0
+# Bounds so a slow or chatty command can never degrade tracing.
+MAX_CUSTOM_TAGS = 20
+MAX_CUSTOM_TAG_LEN = 64
+MAX_SESSION_ID_LEN = 199  # Langfuse session id length limit
 try:
     MAX_CHARS = int(_opt("CC_LANGFUSE_MAX_CHARS") or "20000")
 except ValueError:
@@ -1212,6 +1224,104 @@ def short_session_label(session_id: str, max_len: int = 12) -> str:
 def trace_display_name(session_id: str, turn_num: int) -> str:
     return f"Claude Code - Turn {turn_num} ({short_session_label(session_id)})"
 
+_custom_tags_cache: Optional[List[str]] = None
+_session_label_cache: Optional[str] = None
+
+
+def _run_tag_command(cmd: str, what: str) -> Optional[str]:
+    """Run a user tag/label command, returning stdout or None. Fail-open."""
+    try:
+        proc = subprocess.run(
+            cmd, shell=True, capture_output=True, text=True, timeout=CUSTOM_TAG_TIMEOUT,
+        )
+    except Exception as e:
+        debug(f"{what} command failed: {e!r}")
+        return None
+    if proc.returncode != 0:
+        debug(f"{what} command exited {proc.returncode}: {proc.stderr.strip()[:200]}")
+        return None
+    return proc.stdout
+
+
+def collect_custom_tags() -> List[str]:
+    """Return extra tags from the user-configured CC_LANGFUSE_TAG_COMMAND.
+
+    The command is run once per hook invocation; each non-empty stdout line
+    becomes one tag. This is the supported extension point for attributing
+    traces to whatever a project cares about -- a git branch, a ticket id, a
+    CI run, a cost center -- without editing this hook.
+
+    Fail-open in every direction: no command, a non-zero exit, a timeout, or a
+    crash yields no tags and never interrupts tracing. Output is bounded so a
+    misbehaving command cannot flood a trace with tags.
+    """
+    global _custom_tags_cache
+    if _custom_tags_cache is not None:
+        return _custom_tags_cache
+    _custom_tags_cache = []
+    if not CUSTOM_TAG_COMMAND:
+        return _custom_tags_cache
+    out = _run_tag_command(CUSTOM_TAG_COMMAND, "custom tag")
+    if out is None:
+        return _custom_tags_cache
+    tags: List[str] = []
+    for line in out.splitlines():
+        tag = line.strip()
+        if tag:
+            tags.append(tag[:MAX_CUSTOM_TAG_LEN])
+        if len(tags) >= MAX_CUSTOM_TAGS:
+            break
+    _custom_tags_cache = tags
+    return _custom_tags_cache
+
+
+def resolve_session_label() -> str:
+    """First non-empty stdout line of CC_LANGFUSE_SESSION_LABEL_COMMAND, or ''.
+
+    Cached per process and fail-open. The same script can serve both this and
+    CC_LANGFUSE_TAG_COMMAND: tags use every line, the label uses the first.
+    """
+    global _session_label_cache
+    if _session_label_cache is not None:
+        return _session_label_cache
+    _session_label_cache = ""
+    if not SESSION_LABEL_COMMAND:
+        return _session_label_cache
+    out = _run_tag_command(SESSION_LABEL_COMMAND, "session label")
+    if out is None:
+        return _session_label_cache
+    for line in out.splitlines():
+        label = line.strip()
+        if label:
+            _session_label_cache = label[:MAX_CUSTOM_TAG_LEN]
+            break
+    return _session_label_cache
+
+
+def apply_session_label(session_id: str) -> str:
+    """Group a label's turns into one Langfuse session by relabeling session_id.
+
+    Modes (CC_LANGFUSE_SESSION_LABEL_MODE):
+      prefix   -> "<label>/<session_id>"  (default; one Langfuse session per
+                  Claude session, grouped under the label)
+      collapse -> "<label>"               (every session sharing a label merges
+                  into one Langfuse session)
+      off      -> unchanged
+
+    Returns session_id unchanged when the mode is off or no label resolves.
+    Affects only the trace's grouping id -- the incremental-read state key stays
+    keyed on the raw session id, so relabeling never re-emits past turns.
+    """
+    if SESSION_LABEL_MODE == "off":
+        return session_id
+    label = resolve_session_label()
+    if not label:
+        return session_id
+    if SESSION_LABEL_MODE == "collapse":
+        return label[:MAX_SESSION_ID_LEN]
+    return f"{label}/{session_id}"[:MAX_SESSION_ID_LEN]
+
+
 def get_trace_tags(
     turn: Turn,
     subagent_transcripts_by_tool_use_id: Optional[Dict[str, Dict[str, Any]]] = None,
@@ -1220,7 +1330,10 @@ def get_trace_tags(
     if SKILL_TAGS:
         tags += collect_skill_tags(turn)
         tags += collect_subagent_skill_tags(turn, subagent_transcripts_by_tool_use_id)
-    return tags
+    tags += collect_custom_tags()
+    # De-dup, order-preserving: a custom command may re-emit an existing tag.
+    seen: set = set()
+    return [t for t in tags if not (t in seen or seen.add(t))]
 
 # ---- Generation payloads ----
 def build_generation_input(
@@ -1899,8 +2012,12 @@ def emit_new_turns_from_transcript(
     session_id: str,
     transcript_path: Path,
     *,
+    trace_session_id: Optional[str] = None,
     flush_deferred_agent_turns: bool = False,
 ) -> int:
+    # State/offset tracking keys off the raw session_id; the trace groups under
+    # trace_session_id (which may carry a story label). Defaults to raw.
+    trace_session_id = trace_session_id or session_id
     with FileLock(LOCK_FILE):
         state = load_hook_state()
         key = get_session_state_key(session_id, str(transcript_path))
@@ -1927,7 +2044,7 @@ def emit_new_turns_from_transcript(
         )
         emitted = emit_ready_turns(
             langfuse,
-            session_id,
+            trace_session_id,
             transcript_path,
             turns_to_emit,
             session_state,
@@ -1976,6 +2093,7 @@ def main() -> int:
         return 0
 
     session_id, transcript_path = hook_context
+    trace_session_id = apply_session_label(session_id)
     flush_deferred_agent_turns = is_session_end_hook_payload(payload)
 
     langfuse = create_langfuse_client(config)
@@ -1988,6 +2106,7 @@ def main() -> int:
             config,
             session_id,
             transcript_path,
+            trace_session_id=trace_session_id,
             flush_deferred_agent_turns=flush_deferred_agent_turns,
         )
 
diff --git a/tests/unit/test_custom_tags.py b/tests/unit/test_custom_tags.py
new file mode 100644
index 0000000..580613d
--- /dev/null
+++ b/tests/unit/test_custom_tags.py
@@ -0,0 +1,125 @@
+"""Tests for CC_LANGFUSE_TAG_COMMAND and CC_LANGFUSE_SESSION_LABEL_COMMAND.
+
+Uses the session-scoped ``hook_module`` fixture (see tests/conftest.py). Config is
+read into module globals at import time and results are memoized, so each test
+patches the relevant globals and resets the caches rather than reimporting.
+"""
+from __future__ import annotations
+
+import types
+from typing import Any
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def _reset_caches(hook_module: Any) -> Any:
+    # The caches are process-global; the hook_module fixture is session-scoped.
+    # Reset on both sides so these tests neither inherit nor leak cached state
+    # into other test modules that call get_trace_tags.
+    hook_module._custom_tags_cache = None
+    hook_module._session_label_cache = None
+    yield
+    hook_module._custom_tags_cache = None
+    hook_module._session_label_cache = None
+
+
+def _dummy_turn() -> Any:
+    return types.SimpleNamespace(assistant_msgs=[])
+
+
+# ---- collect_custom_tags ----
+
+def test_two_tags_from_command(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "CUSTOM_TAG_COMMAND", "printf 'sc-1234\\nlayer:build\\n'")
+    assert hook_module.collect_custom_tags() == ["sc-1234", "layer:build"]
+
+
+def test_blank_and_whitespace(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "CUSTOM_TAG_COMMAND", "printf '  sc-9\\n\\n  \\nfoo \\n'")
+    assert hook_module.collect_custom_tags() == ["sc-9", "foo"]
+
+
+def test_no_command(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "CUSTOM_TAG_COMMAND", "")
+    assert hook_module.collect_custom_tags() == []
+
+
+def test_non_zero_exit_fails_open(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "CUSTOM_TAG_COMMAND", "echo nope; exit 3")
+    assert hook_module.collect_custom_tags() == []
+
+
+def test_timeout_fails_open(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "CUSTOM_TAG_COMMAND", "sleep 5; echo late")
+    monkeypatch.setattr(hook_module, "CUSTOM_TAG_TIMEOUT", 1.0)
+    assert hook_module.collect_custom_tags() == []
+
+
+def test_tag_count_cap(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "CUSTOM_TAG_COMMAND", "seq 1 100")
+    assert len(hook_module.collect_custom_tags()) == hook_module.MAX_CUSTOM_TAGS
+
+
+def test_tag_length_cap(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "CUSTOM_TAG_COMMAND", "python3 -c \"print('x'*200)\"")
+    assert len(hook_module.collect_custom_tags()[0]) == hook_module.MAX_CUSTOM_TAG_LEN
+
+
+def test_get_trace_tags_merges_and_dedups(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "SKILL_TAGS", False)
+    monkeypatch.setattr(hook_module, "CUSTOM_TAG_COMMAND", "printf 'claude-code\\nsc-77\\n'")
+    assert hook_module.get_trace_tags(_dummy_turn()) == ["claude-code", "sc-77"]
+
+
+# ---- session label ----
+
+SID = "173624a3-cda5-47d3-97f5-d1b2c68d8fef"
+
+
+def test_session_label_prefix_default(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_COMMAND", "echo sc-53855")
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_MODE", "prefix")
+    assert hook_module.apply_session_label(SID) == f"sc-53855/{SID}"
+
+
+def test_session_label_collapse(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_COMMAND", "echo sc-53855")
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_MODE", "collapse")
+    assert hook_module.apply_session_label(SID) == "sc-53855"
+
+
+def test_session_label_off(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_COMMAND", "echo sc-53855")
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_MODE", "off")
+    assert hook_module.apply_session_label(SID) == SID
+
+
+def test_session_label_no_command(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_COMMAND", "")
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_MODE", "prefix")
+    assert hook_module.apply_session_label(SID) == SID
+
+
+def test_session_label_non_zero_exit_unchanged(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_COMMAND", "echo sc-1; exit 2")
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_MODE", "prefix")
+    assert hook_module.apply_session_label(SID) == SID
+
+
+def test_session_label_capped(hook_module, monkeypatch):
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_COMMAND", "python3 -c \"print('x'*250)\"")
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_MODE", "prefix")
+    out = hook_module.apply_session_label(SID)
+    assert len(out.split("/")[0]) == hook_module.MAX_CUSTOM_TAG_LEN
+    assert len(out) <= hook_module.MAX_SESSION_ID_LEN
+
+
+def test_one_script_feeds_both(hook_module, monkeypatch):
+    script = "printf 'sc-53855\\nlayer:build\\n'"
+    monkeypatch.setattr(hook_module, "SKILL_TAGS", False)
+    monkeypatch.setattr(hook_module, "CUSTOM_TAG_COMMAND", script)
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_COMMAND", script)
+    monkeypatch.setattr(hook_module, "SESSION_LABEL_MODE", "prefix")
+    assert hook_module.get_trace_tags(_dummy_turn()) == ["claude-code", "sc-53855", "layer:build"]
+    assert hook_module.apply_session_label(SID) == f"sc-53855/{SID}"