Newer
Older
navi-1 / navi / core / task_notes.py
"""Pending completion notes for background tasks.

When a background task finishes, the TaskManager records a note here. The
next ``run_stream`` turn drains the notes and injects them into the session
context as a system message (persisted, ``is_display=False``) — the agent
learns the result without the client re-rendering anything (it already saw
the live ``task_update`` event).

Notes are persisted in the shared KV store (scope ``task_notes``) so they
survive a server restart even though the tasks themselves do not.
"""

import asyncio
import json

import structlog

log = structlog.get_logger()

_kv_store = None


def set_kv_store(kv) -> None:
    """Inject the shared KvStore instance (called once at startup)."""
    global _kv_store
    _kv_store = kv


async def add_note(job) -> None:
    """Record a completion note for a finished TaskJob."""
    if _kv_store is None:
        log.warning("tasks.notes_no_store", task_id=job.task_id)
        return
    from navi.config import settings

    async with _lock():
        notes = await _load(job.session_id)
        notes.append({
            "task_id": job.task_id,
            "tool": job.tool,
            "status": job.status,
            "preview": job.preview(limit=800),
            "subagent_tokens": job.subagent_tokens,
        })
        max_pending = settings.task_notes_max_pending
        dropped = 0
        if len(notes) > max_pending:
            dropped = len(notes) - max_pending
            notes = notes[-max_pending:]
        await _save(job.session_id, notes, dropped)


async def drain(session_id: str) -> str | None:
    """Take up to task_notes_per_turn notes and coalesce into one text block."""
    if _kv_store is None:
        return None
    from navi.config import settings

    async with _lock():
        notes = await _load(session_id)
        if not notes:
            return None
        take = notes[: settings.task_notes_per_turn]
        rest = notes[len(take):]
        await _save(session_id, rest, 0)

    lines = [f"- {n['task_id']} ({n['tool']}) {n['status']}: {n['preview'] or '(no output)'}"
             for n in take]
    text = "[Background task results]\n" + "\n".join(lines)
    if rest:
        text += (
            f"\n... {len(rest)} older result(s) not shown — "
            f"call tasks check <task_id> to inspect them."
        )
    text += (
        "\n(Use these results to continue your work. Do not start new "
        "background tasks in response to this note unless the user asked.)"
    )
    return text


_lock_instance = asyncio.Lock()


def _lock() -> asyncio.Lock:
    return _lock_instance


async def _load(session_id: str) -> list[dict]:
    raw = await _kv_store.get("", session_id, "task_notes", "pending")
    if not raw:
        return []
    try:
        data = json.loads(raw)
        return data if isinstance(data, list) else []
    except (ValueError, TypeError):
        return []


async def _save(session_id: str, notes: list[dict], dropped: int) -> None:
    payload = json.dumps(notes, ensure_ascii=False)
    if dropped:
        log.warning("tasks.notes_dropped", session_id=session_id, dropped=dropped)
    await _kv_store.set("", session_id, "task_notes", "pending", payload)


async def pending_count(session_id: str) -> int:
    if _kv_store is None:
        return 0
    return len(await _load(session_id))