✨ Track AgentCore chats with the shared session report writer

lamin_cli/__main__.py

@@ -1238,11 +1238,25 @@ def track_codex_command(name: str | None) -> None:
     return track_codex_session(name=name)
 
 
+@track.command("agentcore")
+@click.option("--name", required=True, help="One-sentence name for this chat's task.")
+def track_agentcore_command(name: str) -> None:
+    """Start or resume tracking inside LaminHub's AgentCore Python runtime."""
+    from lamin_cli.agents.agentcore import track_agentcore_session
+
+    return track_agentcore_session(name=name)
+
+
 def _finish_tracked_session() -> None:
     """Finish a tracked session.
 
-    This can be a shell script run, or a Claude Code, Copilot, Cursor, or Codex session.
+    This can be a shell script or a supported agent session, including AgentCore.
     """
+    from lamin_cli.agents import agentcore
+
+    if agentcore.has_context():
+        return agentcore.finish_agentcore_session()
+
     if os.environ.get("CURSOR_AGENT"):
         from lamin_cli.agents.cursor import finish_cursor_session
         return finish_cursor_session()

lamin_cli/agents/_common.py

@@ -531,7 +531,7 @@ def render_transcript_html(
     entries: list[dict],
     *,
     is_bookkeeping_bash_cmd: Callable[[str], bool],
-    skill_marker: str,
+    skill_marker: str | None,
     shell_tool_names: frozenset[str] = frozenset({"Bash"}),
 ) -> str:
     def content_has_marker(content: object, marker: str) -> bool:
@@ -549,7 +549,8 @@ def content_has_marker(content: object, marker: str) -> bool:
     filtered = [
         msg
         for msg in entries
-        if not content_has_marker(msg.get("content"), skill_marker)
+        if skill_marker is None
+        or not content_has_marker(msg.get("content"), skill_marker)
     ]
 
     bookkeeping_ids: set[str] = {

lamin_cli/agents/agentcore.py

@@ -0,0 +1,160 @@
+"""Track the public conversation supplied by LaminHub's AgentCore runtime.
+
+The runtime refreshes this in-memory snapshot before each Python tool call.
+There is no log discovery, HTTP client or second transcript recorder here.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from uuid import NAMESPACE_URL, UUID, uuid5
+
+import click
+
+from lamin_cli.agents import _common
+
+_TRANSFORM_KEY = "__agentcore__"
+_TRANSFORM_UID = "AgntCoreChat0000"
+
+
+@dataclass(frozen=True)
+class _Context:
+    session_id: str
+    instance_id: str
+    entries: list[dict]
+
+
+_context: _Context | None = None
+
+
+def normalize_messages(messages: list[dict]) -> list[dict]:
+    """Convert Hub's public message strings to the shared renderer's format."""
+    return [
+        {"role": message["role"], "content": [{"type": "text", "text": text}]}
+        for message in messages
+        if message.get("role") in {"user", "assistant"}
+        and isinstance(text := message.get("content"), str)
+        and text.strip()
+    ]
+
+
+def set_context(*, session_id: str, instance_id: str, messages: list[dict]) -> None:
+    """Supply a snapshot from the host, before executing the current tool."""
+    global _context
+    _context = _Context(
+        str(UUID(session_id)), str(UUID(instance_id)), normalize_messages(messages)
+    )
+
+
+def clear_context() -> None:
+    """Release turn-local context along with the worker's credentials."""
+    global _context
+    _context = None
+    os.environ.pop("LAMIN_INITIATED_BY_RUN_UID", None)
+
+
+def has_context() -> bool:
+    return _context is not None
+
+
+def _connected_context(ln) -> _Context:
+    if _context is None:
+        raise click.ClickException(
+            "AgentCore tracking requires a LaminHub runtime context."
+        )
+    if not _common.instance_connected(ln):
+        raise click.ClickException("No LaminDB instance connected.")
+    if str(ln.setup.settings.instance._id) != _context.instance_id:
+        raise click.ClickException("AgentCore context belongs to a different instance.")
+    return _context
+
+
+def _run_uid(context: _Context) -> str:
+    # A stable DB identity survives worker/VM expiry without another state store.
+    return uuid5(
+        NAMESPACE_URL, f"lamin-agentcore:{context.instance_id}:{context.session_id}"
+    ).hex[:20]
+
+
+def get_agentcore_run():
+    """Return this chat's Run, including on a later turn or worker restart."""
+    import lamindb as ln
+
+    context = _connected_context(ln)
+    run = ln.Run.filter(uid=_run_uid(context)).one_or_none()
+    if run is not None and (
+        run.transform.uid != _TRANSFORM_UID
+        or run.reference != context.session_id
+        or run.reference_type != "agentcore"
+    ):
+        raise click.ClickException("AgentCore Run identity does not match this chat.")
+    return run
+
+
+def track_agentcore_session(name: str) -> None:
+    import lamindb as ln
+    from django.db import IntegrityError
+
+    context = _connected_context(ln)
+    run = get_agentcore_run()
+    if run is None:
+        transform, _ = ln.Transform.objects.get_or_create(
+            uid=_TRANSFORM_UID,
+            defaults={
+                "key": _TRANSFORM_KEY,
+                "kind": "function",
+                "description": "An AgentCore chat.",
+            },
+        )
+        run = ln.Run(
+            transform,
+            name=name,
+            status="started",
+            reference=context.session_id,
+            reference_type="agentcore",
+        )
+        run.uid = _run_uid(context)
+        try:
+            run.save()
+        except IntegrityError:
+            # Reconcile a concurrent/retried creation using the DB's unique UID.
+            run = get_agentcore_run()
+            if run is None:
+                raise
+    if run.finished_at is not None:
+        run._status_code = -2  # restarted, as for other harnesses
+        run.finished_at = None
+        run.save()
+    os.environ["LAMIN_INITIATED_BY_RUN_UID"] = run.uid
+    _common.info(f"tracking AgentCore session: {run.uid}")
+
+
+def finish_agentcore_session() -> None:
+    import lamindb as ln
+
+    context = _connected_context(ln)
+    run = get_agentcore_run()
+    if run is None:
+        raise click.ClickException("Run lamin track agentcore --name <task> first.")
+    # The snapshot precedes this tool, so do not wait for a finish command in it.
+    html_doc = _common.render_transcript_html(
+        context.entries,
+        is_bookkeeping_bash_cmd=lambda _: False,
+        skill_marker=None,
+        shell_tool_names=frozenset(),
+    )
+    _common.save_or_replace_report(
+        run,
+        html_doc,
+        context.entries,
+        description="AgentCore session conversation (rendered)",
+        ln=ln,
+    )
+    _common.stamp_transforms(run, context.entries, ln)
+    run._status_code = 0
+    run.finished_at = datetime.now(UTC)
+    run.save()
+    os.environ.pop("LAMIN_INITIATED_BY_RUN_UID", None)
+    _common.info(f"finished tracking AgentCore session: {run.uid}")

tests/agents/test_track_agentcore.py

@@ -0,0 +1,159 @@
+import json
+import os
+from uuid import uuid4
+
+import click
+import lamindb as ln
+import pytest
+from click.testing import CliRunner
+from lamin_cli.__main__ import main
+from lamin_cli.agents import agentcore
+
+
+@pytest.fixture
+def chat(monkeypatch, tmp_path):
+    monkeypatch.chdir(tmp_path)
+    session_id = str(uuid4())
+
+    def supply(messages, *, instance_id=None):
+        agentcore.set_context(
+            session_id=session_id,
+            instance_id=instance_id or str(ln.setup.settings.instance._id),
+            messages=messages,
+        )
+
+    supply([{"role": "user", "content": "Review this dataset"}])
+    yield supply
+    agentcore.clear_context()
+    transform = ln.Transform.filter(uid=agentcore._TRANSFORM_UID).one_or_none()
+    if transform:
+        for run in transform.runs.all():
+            report = run.report
+            run.report = None
+            run.save()
+            run.delete(permanent=True)
+            if report:
+                report.delete(permanent=True)
+        transform.delete(permanent=True)
+
+
+def test_normalize_only_public_hub_text():
+    messages = [
+        {"role": "system", "content": "hidden system"},
+        {"role": "user", "content": "saved text", "access_token": "secret"},
+        {"role": "assistant", "content": "public reply"},
+        {"role": "assistant", "content": " "},
+        {
+            "role": "assistant",
+            "content": [
+                {"reasoningContent": {"reasoningText": {"text": "private reasoning"}}},
+                {
+                    "toolUse": {
+                        "name": "execute_python",
+                        "input": {"code": "hidden tool"},
+                    }
+                },
+            ],
+        },
+    ]
+    entries = agentcore.normalize_messages(messages)
+    assert entries == [
+        {"role": "user", "content": [{"type": "text", "text": "saved text"}]},
+        {"role": "assistant", "content": [{"type": "text", "text": "public reply"}]},
+    ]
+
+
+def test_finish_stores_summary_once_and_reuses_run_after_context_restart(
+    chat, monkeypatch
+):
+    runner = CliRunner()
+    result = runner.invoke(main, ["track", "agentcore", "--name", "Review dataset"])
+    assert result.exit_code == 0, result.output
+    run = agentcore.get_agentcore_run()
+    assert run.status == "started"
+    assert run.uid == os.environ["LAMIN_INITIATED_BY_RUN_UID"]
+    summary = (
+        "Unicode café — <script>alert(1)</script>\n```python\nx = 1\n```\n"
+        + "long summary " * 150
+    )
+    messages = [
+        {"role": "user", "content": "Review this dataset"},
+        {"role": "assistant", "content": summary},
+    ]
+    chat(messages)
+    # Mutation of the caller's data must not change the supplied snapshot.
+    messages[-1]["content"] = "not the snapshot"
+    # Other desktop environment markers must not override the hosted context.
+    monkeypatch.setenv("CODEX_THREAD_ID", "desktop-chat")
+    result = runner.invoke(main, ["finish"])
+    assert result.exit_code == 0, result.output
+    run.refresh_from_db()
+    assert run.status == "completed"
+    assert run.report.n_files == 2
+    entries = [
+        json.loads(line)
+        for line in (run.report.path / "trace.jsonl").read_text().splitlines()
+    ]
+    assert entries[-1]["content"][0]["text"] == summary
+    html = (run.report.path / "index.html").read_text()
+    assert html.count("Unicode café") == 1
+    assert "&lt;script&gt;" in html and "<script>" not in html
+    assert ("long summary " * 150).rstrip() in html
+    assert "not the snapshot" not in html and "lamin finish" not in html
+    report_uid = run.report.uid
+
+    agentcore.clear_context()  # simulate a worker/VM replacement; no local files
+    chat(
+        [
+            {"role": "user", "content": "Follow-up"},
+            {"role": "assistant", "content": "Second summary"},
+        ]
+    )
+    agentcore.track_agentcore_session(name="Follow-up")
+    assert agentcore.get_agentcore_run().uid == run.uid
+    assert agentcore.get_agentcore_run().status == "restarted"
+    agentcore.finish_agentcore_session()
+    agentcore.finish_agentcore_session()  # retry does not create another report
+    run.refresh_from_db()
+    assert run.report.uid == report_uid
+    assert ln.Run.filter(transform=run.transform).count() == 1
+    assert "Second summary" in (run.report.path / "index.html").read_text()
+
+
+def test_finish_failure_leaves_run_open_and_can_retry(chat, monkeypatch):
+    agentcore.track_agentcore_session(name="Retry report")
+    save = agentcore._common.save_or_replace_report
+
+    def fail(*args, **kwargs):
+        raise OSError("storage unavailable")
+
+    monkeypatch.setattr(agentcore._common, "save_or_replace_report", fail)
+    with pytest.raises(OSError, match="storage unavailable"):
+        agentcore.finish_agentcore_session()
+    run = agentcore.get_agentcore_run()
+    assert run.status == "started" and run.report is None
+    monkeypatch.setattr(agentcore._common, "save_or_replace_report", save)
+    agentcore.finish_agentcore_session()
+    assert agentcore.get_agentcore_run().status == "completed"
+
+
+def test_missing_tracking_or_wrong_instance_fails(chat):
+    with pytest.raises(click.ClickException, match="track agentcore"):
+        agentcore.finish_agentcore_session()
+    chat([], instance_id=str(uuid4()))
+    with pytest.raises(click.ClickException, match="different instance"):
+        agentcore.track_agentcore_session(name="wrong instance")
+    agentcore.clear_context()
+    with pytest.raises(click.ClickException, match="runtime context"):
+        agentcore.track_agentcore_session(name="no context")
+
+
+def test_session_and_instance_both_scope_run_identity():
+    session, instance = str(uuid4()), str(uuid4())
+    base = agentcore._Context(session, instance, [])
+    assert agentcore._run_uid(base) != agentcore._run_uid(
+        agentcore._Context(str(uuid4()), instance, [])
+    )
+    assert agentcore._run_uid(base) != agentcore._run_uid(
+        agentcore._Context(session, str(uuid4()), [])
+    )