diff --git a/lamin_cli/__main__.py b/lamin_cli/__main__.py index b5aeb03..cb85c20 100644 --- a/lamin_cli/__main__.py +++ b/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() diff --git a/lamin_cli/agents/_common.py b/lamin_cli/agents/_common.py index 59b7134..846aeda 100644 --- a/lamin_cli/agents/_common.py +++ b/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] = { diff --git a/lamin_cli/agents/agentcore.py b/lamin_cli/agents/agentcore.py new file mode 100644 index 0000000..9702cca --- /dev/null +++ b/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 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}") diff --git a/tests/agents/test_track_agentcore.py b/tests/agents/test_track_agentcore.py new file mode 100644 index 0000000..053c08f --- /dev/null +++ b/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é — \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 "<script>" in html and "