External/roboco
1
0
Fork 0
mirror of https://github.com/rennf93/roboco.git synced 2026-08-03 07:23:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix: open findings cleanup (#122)

Browse Source 
* refactor(usage): remove the unconsumed per-agent USAGE_UPDATE event

USAGE_UPDATE was published per active agent each sweep, bridged, and
broadcast to /ws/system, but no panel client ever consumed it — the
dashboard reads only the aggregate USAGE_SNAPSHOT. Every emission was
wasted event-bus and WebSocket traffic.

Drop the UsageUpdate payload, publish_usage_update and its throttle, the
EventType member, and the bridge subscription. Keep USAGE_SNAPSHOT, which
already carries the per-agent breakdown, so no live data is lost.

* refactor(prompter): remove the legacy local-LLM HTTP endpoints

The panel uses only the live SDK-intake path (/prompter/live/*); the legacy
/prompter/chat, /draft and /sessions/* endpoints — backed by the local Ollama
LLM with hardcoded prompts — had no remaining caller. Remove the router, its
mount in app.py, and its integration test. The live router and the shared
draft-confirmation service are untouched.

* refactor(prompter): drop the dead legacy local-LLM service + schemas

With the legacy HTTP endpoints gone, the local-LLM chat/draft/session methods,
their prompt constants, the ConfirmOverrides/TurnResult dataclasses, and the
entire prompter schema module had no production caller (only their own tests).
Remove them, keeping the live-intake path: create_task_from_draft /
confirm_live_draft, the enum/priority/team coercion, and the pure
description/readiness helpers.

* refactor(agents): stop granting the Task sub-agent tool to roles

Every agent role was granted the built-in Task tool, but no role prompt or
workflow uses it and there are no custom sub-agent definitions — so a Task call
only spawns a context-blind generic sub-agent that burns budget (ToolSearch,
the comment's stated use, is MCP-only and not callable in agent containers).

Drop Task from all three grant points in lockstep: the --tools spawn flag and
both _ROLE_BUILTIN_TOOLS maps (system-prompt + briefing layers), with a
regression guard added to each layer's test.

---------

Co-authored-by: Renn F <rennf93@users.noreply.github.com>
This commit is contained in:
Renzo F
2026-06-12 23:10:21 +02:00
committed by GitHub
co-authored by Renn F
parent cdde3b256e
commit 718d7dd83e
16 changed files with 84 additions and 2941 deletions
Download Patch File Download Diff File Expand all files Collapse all files
roboco/agents/factories/_base.py
+1 -2
View File
@@ -115,14 +115,13 @@ def _autogen_verbs_layer(prompts_path: Path, role: "AgentRole") -> str | None:
# briefing, so we hoist the exact call into a top-of-system-prompt
# layer (highest-priority instruction the model sees).
#
# Read/Bash/Grep/Glob/Task/TodoWrite are needed by every role; Edit/Write
# Read/Bash/Grep/Glob/TodoWrite are needed by every role; Edit/Write
# only by roles that author code or docs.
_BUILTIN_TOOLS_COMMON: tuple[str, ...] = (
"Read",
"Bash",
"Grep",
"Glob",
"Task",
"TodoWrite",
)
_BUILTIN_TOOLS_AUTHORS: tuple[str, ...] = (*_BUILTIN_TOOLS_COMMON, "Edit", "Write")
roboco/api/app.py
-7
View File
@@ -30,7 +30,6 @@ from roboco.api.routes.optimal import router as optimal_router
from roboco.api.routes.orchestrator import router as orchestrator_router
from roboco.api.routes.product import router as product_router
from roboco.api.routes.project import router as project_router
from roboco.api.routes.prompter import router as prompter_router
from roboco.api.routes.prompter_live import router as prompter_live_router
from roboco.api.routes.provider import router as provider_router
from roboco.api.routes.sessions import router as sessions_router
@@ -315,12 +314,6 @@ def create_app() -> FastAPI:
tags=["Providers"],
)
# Prompter — conversational task drafting assistant
app.include_router(
prompter_router,
prefix=f"{api_prefix}/prompter",
tags=["Prompter"],
)
# Prompter live chat — panel <-> spawned intake agent (SSE + relay)
app.include_router(
prompter_live_router,
roboco/api/routes/prompter.py
-298
View File
@@ -1,298 +0,0 @@
"""
Prompter API Routes
Session-based conversational assistant endpoints for drafting tasks:
- POST /api/prompter/sessions : create a new session
- POST /api/prompter/sessions/{id}/messages : send user message, get AI reply
- GET /api/prompter/sessions/{id}/draft : get structured task draft
- POST /api/prompter/sessions/{id}/confirm : confirm draft → create real task
Legacy stateless endpoints (retained for backward compatibility):
- POST /api/prompter/chat : back-and-forth conversation (stateless)
- POST /api/prompter/draft : structured task draft generation (stateless)
"""
from uuid import UUID
from fastapi import APIRouter, HTTPException, status
from roboco.api.deps import CurrentAgentContext, DbSession
from roboco.api.schemas.prompter import (
ChatMessage,
PrompterChatRequest,
PrompterChatResponse,
PrompterDraftRequest,
PrompterDraftResponse,
PrompterDraftTask,
PrompterMessageRequest,
PrompterMessageResponse,
PrompterSessionResponse,
PrompterTurnResponse,
TaskConfirmRequest,
TaskDraftResponse,
)
from roboco.services.base import NotFoundError, ServiceError, ValidationError
from roboco.services.prompter import ConfirmOverrides, get_prompter_service
router = APIRouter()
def _translate_error(e: ServiceError) -> HTTPException:
"""Service errors → HTTP status."""
if isinstance(e, NotFoundError):
return HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail={"error": "not_found", "message": e.message},
)
if isinstance(e, ValidationError):
return HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": "validation_error",
"message": e.message,
"field": e.field,
},
)
return HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": "internal_error", "message": e.message},
)
# =============================================================================
# SESSION-BASED ENDPOINTS
# =============================================================================
@router.post(
"/sessions",
response_model=PrompterSessionResponse,
status_code=status.HTTP_201_CREATED,
)
async def create_session(
db: DbSession,
agent: CurrentAgentContext,
) -> PrompterSessionResponse:
"""Create a new Prompter conversation session linked to the authenticated agent."""
service = get_prompter_service(db)
try:
session = await service.create_session(agent_id=agent.agent_id)
except ServiceError as e:
raise _translate_error(e) from e
# Commit explicitly: the rest of the write surface (tasks, a2a, ...) does
# the same rather than rely on the request-teardown auto-commit, which is
# sensitive to middleware/teardown ordering. Without this the 201 is
# returned but the row may never persist, so the next request 404s.
await db.commit()
return PrompterSessionResponse(
id=UUID(str(session.id)),
agent_id=UUID(str(session.agent_id)),
status=session.status,
created_at=session.created_at,
updated_at=session.updated_at,
)
@router.post(
"/sessions/{session_id}/messages",
response_model=PrompterTurnResponse,
)
async def send_message(
session_id: UUID,
data: PrompterMessageRequest,
db: DbSession,
agent: CurrentAgentContext,
) -> PrompterTurnResponse:
"""
Accept a user message, append it and an AI assistant response to the
conversation, and return the updated message list plus the readiness
signal (``draft_ready`` and the coarse ``scale`` hint) for this turn.
"""
service = get_prompter_service(db)
try:
turn = await service.send_message(
session_id=session_id,
agent_id=agent.agent_id,
content=data.content,
context=data.context,
)
except ServiceError as e:
raise _translate_error(e) from e
await db.commit()
return PrompterTurnResponse(
messages=[
PrompterMessageResponse(
id=UUID(str(msg.id)),
session_id=UUID(str(msg.session_id)),
role=msg.role,
content=msg.content,
created_at=msg.created_at,
)
for msg in turn.messages
],
draft_ready=turn.draft_ready,
scale=turn.scale,
)
@router.get(
"/sessions/{session_id}/draft",
response_model=TaskDraftResponse,
)
async def get_draft(
session_id: UUID,
db: DbSession,
agent: CurrentAgentContext,
) -> TaskDraftResponse:
"""
Return a structured task draft extracted from conversation history via LLM.
The draft contains: title, description, acceptance_criteria, team,
task_type, nature, and estimated_complexity.
"""
service = get_prompter_service(db)
try:
draft_record = await service.get_or_generate_draft(
session_id=session_id,
agent_id=agent.agent_id,
)
except ServiceError as e:
raise _translate_error(e) from e
# Persist a newly generated draft (no-op when it was already cached).
await db.commit()
# Parse the stored draft_data into PrompterDraftTask for validation
try:
draft_task = PrompterDraftTask(**draft_record.draft_data)
except Exception as exc:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={
"error": "draft_schema_error",
"message": f"Stored draft did not match schema: {exc}",
"raw_draft": draft_record.draft_data,
},
) from exc
return TaskDraftResponse(
id=UUID(str(draft_record.id)),
session_id=UUID(str(draft_record.session_id)),
draft=draft_task,
confirmed_at=draft_record.confirmed_at,
task_id=UUID(str(draft_record.task_id)) if draft_record.task_id else None,
created_at=draft_record.created_at,
)
@router.post(
"/sessions/{session_id}/confirm",
response_model=dict,
status_code=status.HTTP_201_CREATED,
)
async def confirm_draft(
session_id: UUID,
data: TaskConfirmRequest,
db: DbSession,
agent: CurrentAgentContext,
) -> dict:
"""
Validate the draft and create a real Task using the existing TaskService.
Returns the created task ID.
"""
service = get_prompter_service(db)
try:
task_id = await service.confirm_draft(
session_id=session_id,
agent_id=agent.agent_id,
confirm_overrides=ConfirmOverrides(
project_id=data.project_id,
product_id=data.product_id,
assigned_to=data.assigned_to,
extra=data.overrides,
draft=data.draft.model_dump(mode="json") if data.draft else None,
),
)
except ServiceError as e:
raise _translate_error(e) from e
await db.commit()
return {"task_id": str(task_id)}
# =============================================================================
# LEGACY STATELESS ENDPOINTS (backward compatibility)
# =============================================================================
@router.post("/chat", response_model=PrompterChatResponse)
async def prompter_chat(
data: PrompterChatRequest,
_agent: CurrentAgentContext,
) -> PrompterChatResponse:
"""
Continue a Prompter conversation (stateless).
The frontend sends the full conversation history (including the new user
message). The assistant replies, optionally signalling that enough context
has been gathered to generate a draft (`draft_ready=True`).
"""
service = get_prompter_service()
try:
result = await service.chat(
messages=[msg.model_dump() for msg in data.messages],
context=data.context,
)
except ServiceError as e:
raise _translate_error(e) from e
return PrompterChatResponse(
message=result["message"],
draft_ready=result["draft_ready"],
)
@router.post("/draft", response_model=PrompterDraftResponse)
async def prompter_draft(
data: PrompterDraftRequest,
_agent: CurrentAgentContext,
) -> PrompterDraftResponse:
"""
Generate a structured task draft from conversation context (stateless).
The frontend sends the full conversation history. The backend calls the
LLM to produce a JSON draft conforming to the TaskCreate schema.
"""
service = get_prompter_service()
try:
result = await service.draft(
messages=[msg.model_dump() for msg in data.messages],
context=data.context,
)
except ServiceError as e:
raise _translate_error(e) from e
draft_raw = result["draft"]
try:
draft = PrompterDraftTask(**draft_raw)
except Exception as e:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={
"error": "draft_schema_error",
"message": f"Generated draft did not match schema: {e}",
"raw_draft": draft_raw,
},
) from e
return PrompterDraftResponse(
draft=draft,
reasoning=result["reasoning"],
)
def _messages_to_dicts(messages: list[ChatMessage]) -> list[dict[str, str]]:
"""Convert ChatMessage list to dict list (internal helper)."""
return [msg.model_dump() for msg in messages]
roboco/api/schemas/prompter.py
-291
View File
@@ -1,291 +0,0 @@
"""
Prompter API Schemas
Request/response models for the conversational Prompter assistant
that helps users draft tasks through natural language.
Includes both the session-based schemas (for the DB-persisted approach)
and the legacy stateless schemas retained for backward compatibility.
"""
from datetime import datetime
from typing import Any
from uuid import UUID
from pydantic import BaseModel, ConfigDict, Field, field_validator
from roboco.models.base import (
Complexity,
TaskNature,
TaskType,
Team,
)
# =============================================================================
# SHARED MESSAGE SCHEMA
# =============================================================================
class ChatMessage(BaseModel):
"""A single message in the Prompter conversation."""
role: str = Field(..., description="One of: user, assistant, system")
content: str = Field(..., min_length=1, description="Message text")
@field_validator("role")
@classmethod
def _valid_role(cls, v: str) -> str:
if v not in {"user", "assistant", "system"}:
raise ValueError("role must be one of: user, assistant, system")
return v
# =============================================================================
# SESSION-BASED SCHEMAS (acceptance-criteria-required names)
# =============================================================================
class PrompterSessionResponse(BaseModel):
"""Response for session creation and retrieval."""
id: UUID
agent_id: UUID
status: str
created_at: datetime
updated_at: datetime | None = None
model_config = ConfigDict(from_attributes=True)
class PrompterMessageRequest(BaseModel):
"""Request body for POST /api/prompter/sessions/{id}/messages."""
content: str = Field(..., min_length=1, description="The user's message text")
context: dict[str, Any] = Field(
default_factory=dict,
description="Optional per-turn context overrides",
)
class PrompterMessageResponse(BaseModel):
"""A single message record returned to the client."""
id: UUID
session_id: UUID
role: str
content: str
created_at: datetime
model_config = ConfigDict(from_attributes=True)
class PrompterTurnResponse(BaseModel):
"""Result of a chat turn: the full message list plus the readiness signal.
Carries ``draft_ready`` (and the coarse ``scale`` hint) so the frontend
consumes the backend's judgement instead of re-deriving it by string match.
"""
messages: list[PrompterMessageResponse]
draft_ready: bool = False
scale: str | None = Field(
default=None,
description="Coarse size hint from the assistant: 'single' or 'multi'",
)
class TaskConfirmRequest(BaseModel):
"""Request body for POST /api/prompter/sessions/{id}/confirm.
Allows the frontend to pass overrides that should be applied
to the draft before the real task is created.
"""
project_id: UUID | None = Field(
default=None,
description="Override project_id from the draft (required if draft omits it)",
)
product_id: UUID | None = Field(
default=None,
description="Override product_id from the draft",
)
assigned_to: str | None = Field(
default=None,
description="Agent slug or UUID to assign the task to",
)
overrides: dict[str, Any] = Field(
default_factory=dict,
description="Additional fields to override in the draft before task creation",
)
draft: "PrompterDraftTask | None" = Field(
default=None,
description=(
"The human-edited structured draft. When present it replaces the "
"stored draft (after re-validation and description re-composition) "
"before the task is created."
),
)
# =============================================================================
# DRAFT TASK SCHEMA (shared between session and legacy paths)
# =============================================================================
class CellWork(BaseModel):
"""One cell's slice of a task's work — the per-cell breakdown of The Work.
For a single-cell task there is exactly one entry; a board-led feature
carries one entry per participating cell.
"""
team: Team = Field(
..., description="The cell (or coordinating team) doing this work"
)
summary: str = Field(
..., min_length=1, description="One-line summary of this cell's slice"
)
items: list[str] = Field(
default_factory=list, description="Concrete deliverables for this cell"
)
class PrompterDraftTask(BaseModel):
"""A task draft produced by the Prompter.
Mirrors TaskCreate fields so the frontend can POST /api/tasks
with confirmed_by_human=True after human review.
The structured spec fields (``objective``, ``what_this_builds``,
``the_work``, ``notes``) are first-class in this contract but persisted
inside the existing ``draft_data`` JSONB column — no migration. The backend
composes ``description`` deterministically from them; ``acceptance_criteria``
renders as Success Criteria.
"""
title: str = Field(..., min_length=1, max_length=200)
description: str = Field(..., min_length=20)
acceptance_criteria: list[str] = Field(..., min_length=1)
team: Team = Field(...)
priority: int = Field(default=2, ge=0, le=3)
task_type: TaskType = Field(...)
nature: TaskNature = Field(...)
estimated_complexity: Complexity = Field(...)
# Structured spec fields — optional for backward compatibility.
objective: str | None = Field(
default=None,
description="The outcome this task delivers, in one or two sentences",
)
what_this_builds: list[str] = Field(
default_factory=list, description="Concrete artifacts this task produces"
)
the_work: list[CellWork] = Field(
default_factory=list,
description="Per-cell breakdown; length drives single vs multi-cell",
)
notes: list[str] = Field(
default_factory=list,
description="Constraints, reuse pointers, things to confirm with the human",
)
project_id: str | None = Field(
default=None,
description=(
"Project UUID as string; exactly one of project_id or "
"product_id must be set"
),
)
product_id: str | None = Field(
default=None,
description=(
"Product UUID as string; exactly one of project_id or "
"product_id must be set"
),
)
assigned_to: str | None = Field(
default=None,
description="Agent slug or UUID to assign the task to",
)
target_date: str | None = Field(
default=None,
description="ISO-8601 target completion date",
)
# Provenance — always set by the prompter backend
source: str = "prompter"
confirmed_by_human: bool = False
# Resolve TaskConfirmRequest.draft now that PrompterDraftTask exists.
TaskConfirmRequest.model_rebuild()
class TaskDraftResponse(BaseModel):
"""Response for GET /api/prompter/sessions/{id}/draft."""
id: UUID
session_id: UUID
draft: PrompterDraftTask
confirmed_at: datetime | None = None
task_id: UUID | None = None
created_at: datetime
model_config = ConfigDict(from_attributes=True)
# =============================================================================
# LEGACY STATELESS SCHEMAS (retained for backward compatibility)
# =============================================================================
class PrompterChatRequest(BaseModel):
"""Request to continue a Prompter conversation (stateless)."""
messages: list[ChatMessage] = Field(
...,
min_length=1,
description="Conversation history including the new user message",
)
context: dict[str, Any] = Field(
default_factory=dict,
description="Optional context (project_id, team, prior drafts, etc.)",
)
class PrompterChatResponse(BaseModel):
"""Response from the Prompter chat endpoint (stateless)."""
message: str = Field(..., description="Assistant's reply")
conversation_id: str | None = Field(
default=None, description="Client-managed conversation identifier"
)
draft_ready: bool = Field(
default=False,
description=(
"True when the assistant believes enough context exists to draft a task"
),
)
class PrompterDraftRequest(BaseModel):
"""Request to generate a task draft from conversation context (stateless)."""
messages: list[ChatMessage] = Field(
..., min_length=1, description="Full conversation used as drafting context"
)
context: dict[str, Any] = Field(
default_factory=dict,
description="Optional overrides (project_id, team, assigned_to, etc.)",
)
class PrompterDraftResponse(BaseModel):
"""Response from the Prompter draft endpoint (stateless)."""
draft: PrompterDraftTask = Field(..., description="Structured task draft")
reasoning: str = Field(
default="",
description="Assistant's explanation of how the draft was derived",
)
model_config = ConfigDict(from_attributes=True)
roboco/api/websocket_bridge.py
+4 -7
View File
@@ -21,7 +21,6 @@ _RATE_LIMIT_WS_TYPES = {
}
_USAGE_WS_TYPES = {
EventType.USAGE_UPDATE: "USAGE_UPDATE",
EventType.USAGE_SNAPSHOT: "USAGE_SNAPSHOT",
}
@@ -150,13 +149,12 @@ async def _handle_rate_limit_event(event: Event) -> None:
async def _handle_usage_event(event: Event) -> None:
"""Forward USAGE_UPDATE/SNAPSHOT events to operator system WS clients.
"""Forward USAGE_SNAPSHOT events to operator system WS clients.
Both event types carry all the fields the panel needs directly in
``event.data``; we tag them with the discriminating ``type`` string the
The event carries all the fields the panel needs directly in
``event.data``; we tag it with the discriminating ``type`` string the
panel switches on (the same UPPER_SNAKE mapping the rate-limit handler
uses), so the panel can distinguish per-agent updates from aggregate
snapshots.
uses).
"""
ws_type = _USAGE_WS_TYPES.get(event.type)
if ws_type is None:
@@ -197,7 +195,6 @@ def register_websocket_bridge_handlers() -> None:
bus.subscribe(EventType.RATE_LIMIT_LIFTED, _handle_rate_limit_event)
# Usage events -> system WebSocket (panel dashboard)
bus.subscribe(EventType.USAGE_UPDATE, _handle_usage_event)
bus.subscribe(EventType.USAGE_SNAPSHOT, _handle_usage_event)
logger.info("WebSocket bridge handlers registered")
roboco/models/events.py
-1
View File
@@ -70,7 +70,6 @@ class EventType(StrEnum):
RATE_LIMIT_LIFTED = "rate_limit.lifted"
# Usage events
USAGE_UPDATE = "usage.update"
USAGE_SNAPSHOT = "usage.snapshot"
# Question events
roboco/runtime/orchestrator.py
+3 -25
View File
@@ -1800,8 +1800,6 @@ class AgentOrchestrator:
- Read/Write/Edit : file IO inside the workspace
- Bash : shell commands (gated by bash-guard hook)
- Grep/Glob : code navigation
- Task : sub-agent dispatch (used for ToolSearch and
other delegated jobs)
- TodoWrite : per-session planning
Permissions still gate *which* paths Edit/Write can touch (see
`_get_role_permissions`), so this is purely about loading vs
@@ -1817,7 +1815,7 @@ class AgentOrchestrator:
"/app/mcp-config.json",
"--strict-mcp-config",
"--tools",
"Read,Write,Edit,Bash,Grep,Glob,Task,TodoWrite",
"Read,Write,Edit,Bash,Grep,Glob,TodoWrite",
"--output-format",
"stream-json",
"--verbose",
@@ -2315,7 +2313,6 @@ class AgentOrchestrator:
"Bash",
"Grep",
"Glob",
"Task",
"TodoWrite",
)
_ROLE_BUILTIN_TOOLS: ClassVar[dict[str, tuple[str, ...]]] = {
@@ -3507,8 +3504,8 @@ class AgentOrchestrator:
current progress without waiting for session close.
Errors per-agent are caught so one bad agent doesn't abort the whole sweep.
Additionally publishes USAGE_UPDATE events per agent (throttled to at most
one per 5-second window) and a USAGE_SNAPSHOT aggregate after the loop.
Also publishes a USAGE_SNAPSHOT aggregate event after the loop so the
/ws/system dashboard updates live for active agents.
"""
if not self._instances:
return
@@ -3548,25 +3545,6 @@ class AgentOrchestrator:
tokens_input, tokens_output = tokens[0], tokens[1]
model = instance.config.model if instance.config else "unknown"
# Publish USAGE_UPDATE event for this agent (throttled).
with contextlib.suppress(Exception):
from roboco.events import get_event_bus
from roboco.services.usage_events import (
UsageUpdate,
publish_usage_update,
)
await publish_usage_update(
get_event_bus(),
UsageUpdate(
agent_id=agent_id,
task_id=instance.current_task_id,
input_tokens=tokens_input,
output_tokens=tokens_output,
model=model,
),
)
# Accumulate per-agent data for the aggregate snapshot.
with contextlib.suppress(Exception):
from roboco.billing.pricing import calculate_cost
roboco/services/prompter.py
+23 -630
View File
@@ -1,13 +1,10 @@
"""
Prompter Service
Conversational LLM assistant that helps users draft tasks.
Uses the project's local LLM (Ollama, OpenAI-compatible) for
natural-language interaction and structured JSON draft generation —
the same engine as RAG/HyDE, so no external API key is required.
Provides both a session-based approach (DB-persisted) and a
legacy stateless interface for backward compatibility.
Server-side helper for the live SDK-intake flow: turns a confirmed structured
draft into a real Task (``create_task_from_draft`` / ``confirm_live_draft``),
plus the pure helpers that compose a task description and parse the interview
readiness signal.
"""
from __future__ import annotations
@@ -16,22 +13,13 @@ import contextlib
import json
import re
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, Literal
from uuid import UUID, uuid4
from uuid import UUID
import httpx
import structlog
from sqlalchemy import select
from roboco.config import settings
from roboco.db.tables import (
AgentTable,
PrompterMessageTable,
PrompterSessionTable,
TaskDraftTable,
TaskTable,
)
from roboco.db.tables import AgentTable, TaskTable
from roboco.foundation.identity import CELL_TEAMS
from roboco.models.base import (
AgentRole,
@@ -42,7 +30,7 @@ from roboco.models.base import (
Team,
)
from roboco.models.task import TaskCreateRequest
from roboco.services.base import NotFoundError, ServiceError, ValidationError
from roboco.services.base import ServiceError, ValidationError
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
@@ -57,22 +45,6 @@ _BOARD_REVIEW_ROLES: frozenset[AgentRole] = frozenset(
)
# ---------------------------------------------------------------------------
# Input types
# ---------------------------------------------------------------------------
@dataclass
class ConfirmOverrides:
"""Optional overrides applied when confirming a draft to create a task."""
project_id: UUID | None = None
product_id: UUID | None = None
assigned_to: str | None = None
extra: dict[str, Any] = field(default_factory=dict)
draft: dict[str, Any] | None = None
@dataclass
class ReadinessTag:
"""Parsed contents of an assistant turn's trailing roboco-meta block."""
@@ -82,144 +54,17 @@ class ReadinessTag:
scale: str | None = None
@dataclass
class TurnResult:
"""Outcome of a chat turn: the message list plus the readiness signal."""
messages: list[PrompterMessageTable]
draft_ready: bool = False
scale: str | None = None
# ---------------------------------------------------------------------------
# Prompts
# ---------------------------------------------------------------------------
_PROMPTER_SYSTEM_PROMPT = (
"You are the RoboCo Prompter — the intake interviewer for an AI agentic "
"software company. A human describes something they want built; you ask a "
"few sharp questions, then a launch-ready task spec is handed to the dev "
"teams.\n\n"
"How RoboCo is organized:\n"
"- A human CEO sits above a Board (Product Owner, Head of Marketing, "
"Auditor).\n"
"- The Main PM coordinates three delivery cells — Backend, Frontend, and "
"UX/UI. Each cell has developers, a QA, a PM, and a documenter.\n"
"- Small, single-domain work (a bug fix, one endpoint, one component) is "
"one task owned by one cell.\n"
"- A real feature is board-led: the Board sets requirements, the Main PM "
"delegates one subtask per participating cell, and the cells deliver in "
"parallel.\n\n"
"What a well-formed task looks like (the house standard):\n"
"- Objective — the outcome, not the implementation.\n"
"- What This Builds — the concrete artifacts.\n"
"- The Work — the per-cell breakdown (one cell for small work; Backend, "
"Frontend, UX/UI for a feature).\n"
"- Notes — constraints, what to reuse, anything to confirm with the human.\n"
"- Success Criteria — verifiable acceptance criteria.\n\n"
"Your interview discipline:\n"
"- Open by reflecting back, in one or two sentences, what you understand "
"they want, so they can correct course immediately.\n"
"- Then ask only the highest-leverage questions you are actually missing — "
"one or two per turn. Never dump a checklist.\n"
"- Before you can draft, cover: (1) the true objective, (2) scope "
"boundaries — what is explicitly out, (3) the surface — which page, "
"endpoint, or component, grounded in the projects/products you are shown, "
"(4) reuse vs build — what existing code or services to lean on, "
"(5) the audience, (6) what 'done' looks like.\n"
"- Stop as soon as objective, scope, surface, and acceptance are clear. "
"Aim for two to four turns total. Do not pad the conversation.\n"
"- Use the real project and product names you are given; prefer an "
"existing surface over inventing one.\n\n"
"Every reply ends with exactly one fenced control block the human never "
"sees, reporting coverage and readiness:\n"
"```roboco-meta\n"
'{"covered": ["objective", "scope", "surface", "acceptance"], '
'"ready": false, "scale": "single"}\n'
"```\n"
"- covered: which of objective / scope / surface / reuse / audience / "
"acceptance you have nailed down.\n"
"- ready: true only when you could write a complete task spec right now.\n"
"- scale: 'single' for one-cell work, 'multi' for a board-led feature "
"across cells.\n"
"Write nothing after that block."
)
_DRAFT_SYSTEM_PROMPT = (
"You are the RoboCo Prompter's drafting engine. Given a finished "
"conversation, output a single JSON object — a structured task "
"draft. No markdown, no prose, no code fence.\n\n"
"Required fields:\n"
"- title: concise, actionable (max 200 chars).\n"
"- objective: the outcome in one or two sentences.\n"
"- what_this_builds: array of concrete artifacts (strings).\n"
"- the_work: array of per-cell slices. Each item is "
'{"team": backend|frontend|ux_ui, "summary": one line, '
'"items": [deliverables]}. One entry for single-cell work; one entry per '
"participating cell for a board-led feature.\n"
"- acceptance_criteria: array of verifiable criteria (at least one) — "
"these become Success Criteria.\n"
"- notes: array of constraints, reuse pointers, things to confirm (may be "
"empty).\n"
"- team: the primary cell (backend|frontend|ux_ui). For a multi-cell "
"feature set the lead cell here; the backend routes it through the Main "
"PM.\n"
"- task_type: one of code, documentation, research, planning, design, "
"administrative.\n"
"- nature: technical or non_technical.\n"
"- estimated_complexity: low, medium, high.\n"
"- priority: integer 0-3 (0 highest, 3 lowest).\n\n"
"Optional, only if unambiguous from context: project_id, product_id, "
"assigned_to, target_date.\n\n"
"Do NOT write a 'description' field — the backend composes it from your "
"structured fields.\n\n"
"Example shape (abbreviated):\n"
'{"title": "...", "objective": "...", "what_this_builds": ["..."], '
'"the_work": [{"team": "backend", "summary": "...", "items": ["..."]}, '
'{"team": "frontend", "summary": "...", "items": ["..."]}], '
'"acceptance_criteria": ["..."], "notes": ["..."], "team": "backend", '
'"task_type": "code", "nature": "technical", '
'"estimated_complexity": "high", "priority": 1}\n\n'
"Return ONLY the JSON object."
)
class PrompterService:
"""Service for Prompter chat, session management, and structured draft generation.
"""Create tasks from confirmed intake drafts.
Accepts an optional SQLAlchemy ``AsyncSession`` for the session-based
(DB-persisted) interface. When no session is provided, only the legacy
stateless ``chat()`` and ``draft()`` methods are available.
Accepts an optional SQLAlchemy ``AsyncSession`` for the DB-backed task
creation. The pure draft/description helpers below need no session.
"""
def __init__(self, db: AsyncSession | None = None) -> None:
self.log = logger.bind(component="prompter_service")
self._db = db
async def _create_message(
self, *, messages: list[dict[str, str]], max_tokens: int
) -> str:
"""Call the local LLM and return the reply text.
Uses the project's local LLM — the same OpenAI-compatible Ollama
endpoint as RAG/HyDE (``settings.local_llm_*``), so no external API key
is required. ``messages`` is an OpenAI-style list (system + turns). This
is the single seam the prompter tests substitute.
"""
async with httpx.AsyncClient(timeout=120.0) as client:
resp = await client.post(
f"{settings.local_llm_base_url}/chat/completions",
json={
"model": settings.local_llm_model,
"messages": messages,
"max_tokens": max_tokens,
"options": {"num_ctx": 8192},
},
)
resp.raise_for_status()
data = resp.json()
return str(data["choices"][0]["message"]["content"] or "").strip()
@property
def _session(self) -> AsyncSession:
"""Return DB session, raising if not configured."""
@@ -230,180 +75,6 @@ class PrompterService:
)
return self._db
# -----------------------------------------------------------------------
# Session-based interface
# -----------------------------------------------------------------------
async def create_session(self, agent_id: UUID) -> PrompterSessionTable:
"""Create a new Prompter conversation session."""
session = PrompterSessionTable(
id=uuid4(),
agent_id=agent_id,
status="active",
created_at=datetime.now(UTC),
)
self._session.add(session)
await self._session.flush()
self.log.info("Prompter session created", session_id=str(session.id))
return session
async def send_message(
self,
session_id: UUID,
agent_id: UUID,
content: str,
context: dict[str, Any] | None = None,
) -> TurnResult:
"""
Append a user message, call the LLM for a reply, persist both, and
return all messages plus the readiness signal for this turn.
"""
session = await self._get_session(session_id, agent_id)
# Persist the user message first
user_msg = PrompterMessageTable(
id=uuid4(),
session_id=session_id,
role="user",
content=content,
created_at=datetime.now(UTC),
)
self._session.add(user_msg)
await self._session.flush()
# Load full conversation history for the LLM call
history = await self._load_messages(session_id)
chat_messages = [{"role": m.role, "content": m.content} for m in history]
# Ground the interview in the real projects/products the human can target
live_context = await self._assemble_live_context()
# Call the LLM
llm_reply = await self._llm_chat(
messages=chat_messages,
context=context,
live_context=live_context,
)
# Persist the assistant reply (control block already stripped)
assistant_msg = PrompterMessageTable(
id=uuid4(),
session_id=session_id,
role="assistant",
content=llm_reply["message"],
created_at=datetime.now(UTC),
)
self._session.add(assistant_msg)
# Update session status if draft is ready
if llm_reply["draft_ready"] and session.status == "active":
session.status = "draft_ready"
await self._session.flush()
self.log.info(
"Message processed",
session_id=str(session_id),
draft_ready=llm_reply["draft_ready"],
)
return TurnResult(
messages=await self._load_messages(session_id),
draft_ready=bool(llm_reply["draft_ready"]),
scale=llm_reply.get("scale"),
)
async def get_or_generate_draft(
self,
session_id: UUID,
agent_id: UUID,
) -> TaskDraftTable:
"""
Return an existing draft for the session, or generate one via LLM
if none exists yet.
"""
await self._get_session(session_id, agent_id)
# Check for an existing draft
result = await self._session.execute(
select(TaskDraftTable)
.where(TaskDraftTable.session_id == session_id)
.order_by(TaskDraftTable.created_at.desc())
.limit(1)
)
existing = result.scalar_one_or_none()
if existing is not None:
return existing
# No draft yet — generate one from conversation history
history = await self._load_messages(session_id)
if not history:
raise ValidationError(
message=(
"Cannot generate a draft from an empty conversation; "
"send at least one message first."
),
field="messages",
)
chat_messages = [{"role": m.role, "content": m.content} for m in history]
draft_result = await self._llm_draft(
messages=chat_messages,
)
draft_record = TaskDraftTable(
id=uuid4(),
session_id=session_id,
draft_data=draft_result["draft"],
created_at=datetime.now(UTC),
)
self._session.add(draft_record)
await self._session.flush()
return draft_record
async def confirm_draft(
self,
session_id: UUID,
agent_id: UUID,
confirm_overrides: ConfirmOverrides | None = None,
) -> UUID:
"""
Validate the draft and create a real Task via the TaskService.
Returns the newly created task's UUID.
"""
session_rec = await self._get_session(session_id, agent_id)
ov = confirm_overrides or ConfirmOverrides()
# Get or generate the draft. A human-edited structured draft, if passed,
# replaces the stored one before overrides and re-composition.
draft_record = await self.get_or_generate_draft(session_id, agent_id)
if ov.draft is not None:
draft_data: dict[str, Any] = dict(ov.draft)
draft_data["source"] = "prompter"
draft_data["confirmed_by_human"] = False
else:
draft_data = dict(draft_record.draft_data)
self._apply_overrides(draft_data, ov)
task = await self.create_task_from_draft(draft_data, agent_id)
# Persist the launched draft so the stored record reflects reality.
draft_record.draft_data = draft_data
# Mark draft as confirmed
now = datetime.now(UTC)
draft_record.confirmed_at = now
draft_record.task_id = task.id
session_rec.status = "confirmed"
await self._session.flush()
self.log.info(
"Draft confirmed — task created",
session_id=str(session_id),
task_id=str(task.id),
)
return UUID(str(task.id))
async def _assignee_is_board(self, agent_id: UUID) -> bool:
"""True if ``agent_id`` is a board/advisory role (PO / marketing / auditor)."""
result = await self._session.execute(
@@ -421,18 +92,16 @@ class PrompterService:
) -> TaskTable:
"""Create a Task from a structured draft.
Shared by both prompter confirm paths (``confirm_draft`` and the
live-intake ``confirm_live_draft``): recomposes the description,
validates exactly-one target, coerces enums, routes the owning team
(product → Main PM, project → lead cell), and persists via
``TaskService.create``. Mutates ``draft_data['description']`` in place.
``confirmed_by_human=True`` — the CEO confirmed it.
Recomposes the description, validates exactly-one target, coerces enums,
routes the owning team (product → Main PM, project → lead cell), and
persists via ``TaskService.create``. Mutates ``draft_data['description']``
in place. ``confirmed_by_human=True`` — the CEO confirmed it.
``status`` defaults to ``BACKLOG`` (legacy ``confirm_draft`` behaviour).
The live-intake buttons pass ``PENDING`` + an ``assigned_to`` (a board
agent for "Board review & Start", main-pm for "Approve & Start") so the
task starts immediately on the chosen review path. An explicit
``assigned_to`` wins over any assignee carried on the draft.
``status`` defaults to ``BACKLOG``. The live-intake buttons pass
``PENDING`` + an ``assigned_to`` (a board agent for "Board review &
Start", main-pm for "Approve & Start") so the task starts immediately on
the chosen review path. An explicit ``assigned_to`` wins over any
assignee carried on the draft.
"""
# Recompose the description from the (possibly edited) structured fields —
# the task always carries a freshly-composed, consistent description.
@@ -552,18 +221,6 @@ class PrompterService:
)
return UUID(str(task.id))
@staticmethod
def _apply_overrides(draft_data: dict[str, Any], ov: ConfirmOverrides) -> None:
"""Merge confirm-time overrides onto the draft data in place."""
if ov.project_id is not None:
draft_data["project_id"] = str(ov.project_id)
if ov.product_id is not None:
draft_data["product_id"] = str(ov.product_id)
if ov.assigned_to is not None:
draft_data["assigned_to"] = ov.assigned_to
if ov.extra:
draft_data.update(ov.extra)
@staticmethod
def _resolve_uuid_field(draft_data: dict[str, Any], key: str) -> UUID | None:
"""Parse ``draft_data[key]`` as a UUID; None if absent, raises if malformed."""
@@ -653,266 +310,16 @@ class PrompterService:
return 2
return 2
# -----------------------------------------------------------------------
# Private helpers (session-based)
# -----------------------------------------------------------------------
async def _get_session(
self, session_id: UUID, agent_id: UUID
) -> PrompterSessionTable:
"""Load and authorize a PrompterSession."""
result = await self._session.execute(
select(PrompterSessionTable).where(PrompterSessionTable.id == session_id)
)
rec = result.scalar_one_or_none()
if rec is None:
raise NotFoundError("Prompter session", str(session_id))
if rec.agent_id != agent_id:
raise ServiceError(
f"Session {session_id} does not belong to agent {agent_id}"
)
return rec
async def _load_messages(self, session_id: UUID) -> list[PrompterMessageTable]:
"""Return all messages for a session ordered by creation time."""
result = await self._session.execute(
select(PrompterMessageTable)
.where(PrompterMessageTable.session_id == session_id)
.order_by(PrompterMessageTable.created_at)
)
return list(result.scalars().all())
async def _assemble_live_context(self) -> str | None:
"""Build a compact 'Available projects / products' block for the interview.
Grounds the assistant in the real targets the human can launch against,
so it references existing surfaces and can resolve project/product
itself. Best-effort: a lookup failure degrades to no context rather than
breaking the chat. Returns None when nothing is registered.
"""
if self._db is None:
return None
from roboco.services.product import get_product_service
from roboco.services.project import get_project_service
lines: list[str] = []
try:
projects = await get_project_service(self._session).list_all(
active_only=True, limit=50
)
except Exception as exc:
self.log.warning("Live project list unavailable", error=str(exc))
projects = []
if projects:
lines.append("Available projects (single-cell tasks target one of these):")
lines.extend(f" - {p.name} (slug: {p.slug}, id: {p.id})" for p in projects)
try:
products = await get_product_service(self._session).list_all(limit=50)
except Exception as exc:
self.log.warning("Live product list unavailable", error=str(exc))
products = []
if products:
lines.append(
"Available products (board-led multi-cell features target one "
"of these):"
)
lines.extend(
f" - {pr.name} (slug: {pr.slug}, id: {pr.id})" for pr in products
)
return "\n".join(lines) if lines else None
# -----------------------------------------------------------------------
# Shared LLM helpers
# -----------------------------------------------------------------------
async def _llm_chat(
self,
messages: list[dict[str, str]],
context: dict[str, Any] | None = None,
max_tokens: int = 2048,
live_context: str | None = None,
) -> dict[str, Any]:
"""Call the LLM for a chat response.
Returns ``{message, draft_ready, scale}`` where ``message`` is the
user-visible reply with the trailing roboco-meta control block stripped.
"""
user_prompt = _build_chat_prompt(messages, context, live_context)
try:
content = await self._create_message(
messages=[
{"role": "system", "content": _PROMPTER_SYSTEM_PROMPT},
{"role": "user", "content": user_prompt},
],
max_tokens=max_tokens,
)
except Exception as e:
self.log.error("Prompter chat LLM call failed", error=str(e))
raise ServiceError(f"LLM chat failed: {e}") from e
if not content:
raise ServiceError("LLM returned empty content")
clean, tag = parse_readiness(content)
# If the model omitted the control block, fall back to the clean text
# so the user still sees a reply rather than an empty bubble.
message = clean or content
return {
"message": message,
"draft_ready": bool(tag and tag.ready),
"scale": tag.scale if tag else None,
}
async def _llm_draft(
self,
messages: list[dict[str, str]],
context: dict[str, Any] | None = None,
max_tokens: int = 4096,
) -> dict[str, Any]:
"""Call the LLM to generate a structured draft. Returns {draft, reasoning}."""
user_prompt = _build_draft_prompt(messages, context)
try:
content = await self._create_message(
messages=[
{"role": "system", "content": _DRAFT_SYSTEM_PROMPT},
{"role": "user", "content": user_prompt},
],
max_tokens=max_tokens,
)
except Exception as e:
self.log.error("Prompter draft LLM call failed", error=str(e))
raise ServiceError(f"LLM draft generation failed: {e}") from e
if not content:
raise ServiceError("LLM returned empty content for draft")
try:
draft_data = json.loads(_strip_code_fences(content))
except json.JSONDecodeError as e:
self.log.warning("Draft JSON parse failed", content_preview=content[:200])
raise ValidationError(
message=f"Draft response was not valid JSON: {e}",
field="draft",
) from e
draft_data["source"] = "prompter"
draft_data["confirmed_by_human"] = False
# Compose the markdown description from the structured fields — the model
# never hand-formats it, so the description is always consistent.
draft_data["description"] = compose_description(draft_data)
return {
"draft": draft_data,
"reasoning": _build_reasoning(messages, draft_data),
}
# -----------------------------------------------------------------------
# Legacy stateless interface
# -----------------------------------------------------------------------
async def chat(
self,
messages: list[dict[str, str]],
context: dict[str, Any] | None = None,
max_tokens: int = 2048,
) -> dict[str, Any]:
"""Continue a Prompter conversation (stateless)."""
return await self._llm_chat(
messages=messages,
context=context,
max_tokens=max_tokens,
)
async def draft(
self,
messages: list[dict[str, str]],
context: dict[str, Any] | None = None,
max_tokens: int = 4096,
) -> dict[str, Any]:
"""Generate a structured task draft from conversation context (stateless)."""
return await self._llm_draft(
messages=messages,
context=context,
max_tokens=max_tokens,
)
# ---------------------------------------------------------------------------
# Module-level helpers (pure functions, no state)
# ---------------------------------------------------------------------------
def _build_chat_prompt(
messages: list[dict[str, str]],
context: dict[str, Any] | None,
live_context: str | None = None,
) -> str:
lines: list[str] = []
if live_context:
lines.append(live_context)
lines.append("")
if context:
lines.append("Context:")
for key, value in context.items():
lines.append(f" {key}: {value}")
lines.append("")
lines.append("Conversation:")
for msg in messages:
role = msg.get("role", "user")
content = msg.get("content", "")
lines.append(f"{role}: {content}")
lines.append("")
lines.append(
"Continue the conversation as the Prompter assistant. End with the "
"roboco-meta control block. If you can write a complete task spec now, "
"set ready to true."
)
return "\n".join(lines)
def _build_draft_prompt(
messages: list[dict[str, str]],
context: dict[str, Any] | None,
) -> str:
lines: list[str] = []
lines.append(
"Produce a JSON task draft from the following conversation. "
"Return ONLY valid JSON — no markdown, no preamble."
)
if context:
lines.append("")
lines.append("Overrides:")
for key, value in context.items():
lines.append(f" {key}: {value}")
lines.append("")
lines.append("Conversation:")
for msg in messages:
role = msg.get("role", "user")
content = msg.get("content", "")
lines.append(f"{role}: {content}")
return "\n".join(lines)
def _strip_code_fences(content: str) -> str:
"""Strip a wrapping markdown code fence (```json ... ```) if present.
Local models often wrap JSON output in a fenced block; drop the opening
fence line and the closing fence so the body parses cleanly as JSON.
"""
text = content.strip()
if text.startswith("```"):
text = text.split("\n", 1)[1] if "\n" in text else text[3:]
if text.rstrip().endswith("```"):
text = text.rstrip()[:-3]
return text.strip()
_META_FENCE_RE = re.compile(r"```roboco-meta\s*(.*?)```", re.DOTALL)
# Mirror of PrompterDraftTask.description min_length — below this the composed
# body is too thin to be a valid task, so we fall back to any provided text.
# Below this length the composed body is too thin to be a valid task, so we
# fall back to any model-provided description text.
_MIN_DESCRIPTION_LEN = 20
_TEAM_LABELS: dict[str, str] = {
@@ -1050,20 +457,6 @@ def compose_description(draft: dict[str, Any]) -> str:
return _text(draft.get("description")) or composed
def _build_reasoning(
messages: list[dict[str, str]],
draft_data: dict[str, Any],
) -> str:
title = draft_data.get("title", "Untitled")
team = draft_data.get("team", "unknown")
complexity = draft_data.get("estimated_complexity", "unknown")
return (
f"Draft generated from conversation of {len(messages)} messages. "
f"Proposed task '{title}' for team {team} "
f"with complexity {complexity}."
)
# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------
@@ -1072,7 +465,7 @@ def _build_reasoning(
def get_prompter_service(db: AsyncSession | None = None) -> PrompterService:
"""Create a PrompterService instance.
Pass ``db`` for the session-based interface; omit for the stateless
legacy interface.
Pass ``db`` for the DB-backed task-creation interface; omit for the pure
draft/description helpers.
"""
return PrompterService(db=db)
roboco/services/usage_events.py
+4 -82
View File
@@ -1,21 +1,15 @@
"""
Usage Event Publisher
Throttled helpers for publishing USAGE_UPDATE and USAGE_SNAPSHOT events
to the StreamEventBus. Consumed by the orchestrator token sweep and
forwarded to /ws/system WebSocket clients via the websocket_bridge.
Helper for publishing USAGE_SNAPSHOT aggregate events to the StreamEventBus.
Consumed by the orchestrator token sweep and forwarded to /ws/system WebSocket
clients via the websocket_bridge.
Throttle window: one USAGE_UPDATE publish per agent per 5-second window.
Subsequent calls within the window are silently dropped so a frequent
sweep loop cannot flood the event bus or WebSocket clients.
USAGE_SNAPSHOT is always published (no per-agent throttle) — it is an
aggregate and published at most once per sweep cycle.
USAGE_SNAPSHOT is an aggregate, published at most once per sweep cycle.
"""
from __future__ import annotations
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any
@@ -23,64 +17,6 @@ from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from roboco.events.stream_bus import StreamEventBus
_THROTTLE_WINDOW_SECONDS: float = 5.0
class _UsageThrottle:
"""Per-agent last-publish timestamp tracker.
Uses ``time.monotonic()`` so clock adjustments (NTP, DST) don't cause
spurious suppressions or double-fires.
"""
def __init__(self, window: float = _THROTTLE_WINDOW_SECONDS) -> None:
self._window = window
self._last: dict[str, float] = {}
def should_publish(self, agent_id: str) -> bool:
"""Return True if the agent is outside the throttle window.
Also records the current monotonic time as the new *last published*
timestamp when it returns True, so the caller does not need to call a
separate ``record()`` method.
"""
now = time.monotonic()
if now - self._last.get(agent_id, 0.0) >= self._window:
self._last[agent_id] = now
return True
return False
# Module-level singleton — shared across all callers in this process.
_throttle = _UsageThrottle()
@dataclass(frozen=True, slots=True)
class UsageUpdate:
"""Per-agent cumulative token counts carried by a USAGE_UPDATE event.
Fields map directly onto the event payload the panel consumes (the
backend emits these per active agent during the token sweep).
"""
agent_id: str
task_id: str | None
input_tokens: int
output_tokens: int
model: str
timestamp: datetime | None = None
def event_data(self) -> dict[str, Any]:
"""Render the event payload, stamping ``timestamp`` if not supplied."""
return {
"agent_id": self.agent_id,
"task_id": self.task_id,
"input_tokens": self.input_tokens,
"output_tokens": self.output_tokens,
"model": self.model,
"timestamp": (self.timestamp or datetime.now(UTC)).isoformat(),
}
@dataclass(frozen=True, slots=True)
class UsageSnapshot:
@@ -108,20 +44,6 @@ class UsageSnapshot:
}
async def publish_usage_update(bus: StreamEventBus, update: UsageUpdate) -> bool:
"""Publish a USAGE_UPDATE event if the per-agent throttle window has elapsed.
Returns True if the event was published; False if suppressed by the throttle.
"""
if not _throttle.should_publish(update.agent_id):
return False
from roboco.models.events import Event, EventType # lazy — avoids circular import
await bus.publish(Event(type=EventType.USAGE_UPDATE, data=update.event_data()))
return True
async def publish_usage_snapshot(bus: StreamEventBus, snapshot: UsageSnapshot) -> None:
"""Publish a USAGE_SNAPSHOT aggregate event (no throttle)."""
from roboco.models.events import Event, EventType # lazy — avoids circular import
tests/integration/test_prompter_routes.py
-862
View File
@@ -1,862 +0,0 @@
"""Prompter API route integration tests.
Covers both the new session-based endpoints:
POST /sessions, POST /sessions/{id}/messages, GET /sessions/{id}/draft,
POST /sessions/{id}/confirm
And the legacy stateless endpoints:
POST /chat, POST /draft
"""
from __future__ import annotations
import json
from http import HTTPStatus
from typing import TYPE_CHECKING, Any
from unittest.mock import AsyncMock, patch
from uuid import uuid4
import pytest
import pytest_asyncio
from fastapi import FastAPI
from httpx import ASGITransport, AsyncClient
from roboco.api.deps import get_agent_context, get_db
from roboco.api.routes.prompter import router as prompter_router
from roboco.db.tables import AgentTable, ProjectTable
from roboco.models.base import AgentRole, AgentStatus, Team
from roboco.models.permissions import AgentContext
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
if TYPE_CHECKING:
from collections.abc import AsyncIterator
from sqlalchemy.ext.asyncio import AsyncSession
# Expected message counts in multi-turn tests
_SINGLE_TURN_MSGS = 2 # 1 user + 1 assistant
_DOUBLE_TURN_MSGS = 4 # 2 user + 2 assistant
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest_asyncio.fixture
async def prompter_client(
db_session: AsyncSession,
) -> AsyncIterator[dict[str, Any]]:
agent_id = uuid4()
agent = AgentTable(
id=agent_id,
name="DevAgent",
slug=f"dev-agent-{uuid4().hex[:8]}",
role=AgentRole.DEVELOPER,
team=None,
status=AgentStatus.ACTIVE,
model_config={},
system_prompt="dev",
capabilities=[],
permissions={},
metrics={},
)
db_session.add(agent)
await db_session.flush()
app = FastAPI()
app.include_router(prompter_router, prefix="/api/prompter")
async def _override_db() -> AsyncIterator[AsyncSession]:
yield db_session
async def _override_agent() -> AgentContext:
return AgentContext(
agent_id=agent_id,
role=AgentRole.DEVELOPER,
team=None,
)
app.dependency_overrides[get_db] = _override_db
app.dependency_overrides[get_agent_context] = _override_agent
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as client:
yield {"client": client, "agent": agent, "db": db_session}
app.dependency_overrides.clear()
@pytest_asyncio.fixture
async def project_fixture(db_session: AsyncSession) -> ProjectTable:
"""Create a minimal project for task creation in confirm tests."""
creator = AgentTable(
id=uuid4(),
name="ProjectCreator",
slug=f"proj-creator-{uuid4().hex[:8]}",
role=AgentRole.MAIN_PM,
team=Team.BACKEND,
status=AgentStatus.ACTIVE,
model_config={},
system_prompt="pm",
capabilities=[],
permissions={},
metrics={},
)
db_session.add(creator)
await db_session.flush()
project = ProjectTable(
id=uuid4(),
name="Test Project",
slug=f"test-project-{uuid4().hex[:8]}",
git_url="https://github.com/test/repo.git",
default_branch="main",
assigned_cell=Team.BACKEND,
created_by=creator.id,
)
db_session.add(project)
await db_session.flush()
return project
_HDR = {"X-Agent-ID": "be-dev-1", "X-Agent-Role": "developer"}
@pytest_asyncio.fixture
async def cross_request_client(
_test_database_url: str,
) -> AsyncIterator[dict[str, Any]]:
"""Client whose DB dependency yields a fresh, NON-auto-committing session
per request.
This is the boundary the shared-session ``prompter_client`` fixture can't
exercise: here a write is only visible to the next request if the route
committed it explicitly. The seed agent is committed up front so both
requests can resolve it.
"""
engine = create_async_engine(_test_database_url, future=True, pool_pre_ping=True)
maker = async_sessionmaker(bind=engine, expire_on_commit=False)
agent_id = uuid4()
async with maker() as seed:
seed.add(
AgentTable(
id=agent_id,
name="XReqAgent",
slug=f"xreq-agent-{uuid4().hex[:8]}",
role=AgentRole.DEVELOPER,
team=None,
status=AgentStatus.ACTIVE,
model_config={},
system_prompt="dev",
capabilities=[],
permissions={},
metrics={},
)
)
await seed.commit()
app = FastAPI()
app.include_router(prompter_router, prefix="/api/prompter")
async def _override_db() -> AsyncIterator[AsyncSession]:
# A fresh session per request that does NOT commit on teardown, so
# persistence depends solely on the route's explicit commit.
async with maker() as session:
yield session
async def _override_agent() -> AgentContext:
return AgentContext(agent_id=agent_id, role=AgentRole.DEVELOPER, team=None)
app.dependency_overrides[get_db] = _override_db
app.dependency_overrides[get_agent_context] = _override_agent
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as client:
yield {"client": client, "agent_id": agent_id}
app.dependency_overrides.clear()
await engine.dispose()
@pytest.mark.asyncio
async def test_session_persists_across_requests(cross_request_client: dict) -> None:
"""A created session must survive into the next request's own DB session.
Regression for the production 404: the create returned 201 but the session
write was never committed, so the immediately-following /messages call could
not find it. Without the route's explicit commit, this is a 404.
"""
client = cross_request_client["client"]
create = await client.post("/api/prompter/sessions", json={}, headers=_HDR)
assert create.status_code == HTTPStatus.CREATED
session_id = create.json()["id"]
reply = (
'ack\n```roboco-meta\n{"covered": [], "ready": false, "scale": "single"}\n```'
)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=reply,
):
msg = await client.post(
f"/api/prompter/sessions/{session_id}/messages",
json={"content": "hello"},
headers=_HDR,
)
assert msg.status_code == HTTPStatus.OK, msg.json()
assert len(msg.json()["messages"]) == _SINGLE_TURN_MSGS
# =============================================================================
# Session-based endpoint tests
# =============================================================================
@pytest.mark.asyncio
async def test_create_session_success(prompter_client: dict) -> None:
"""POST /sessions creates a new session linked to the agent."""
client = prompter_client["client"]
response = await client.post(
"/api/prompter/sessions",
json={},
headers=_HDR,
)
assert response.status_code == HTTPStatus.CREATED
body = response.json()
assert "id" in body
assert body["status"] == "active"
assert "agent_id" in body
assert "created_at" in body
@pytest.mark.asyncio
async def test_send_message_success(prompter_client: dict) -> None:
"""POST /sessions/{id}/messages appends user+assistant messages."""
client = prompter_client["client"]
# Create session
session_resp = await client.post("/api/prompter/sessions", json={}, headers=_HDR)
session_id = session_resp.json()["id"]
mock_response = "Great! Let's gather requirements."
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=mock_response,
):
response = await client.post(
f"/api/prompter/sessions/{session_id}/messages",
json={"content": "I need a new feature"},
headers=_HDR,
)
assert response.status_code == HTTPStatus.OK
body = response.json()
messages = body["messages"]
assert len(messages) == _SINGLE_TURN_MSGS
roles = [m["role"] for m in messages]
assert "user" in roles
assert "assistant" in roles
assert messages[-1]["content"] == "Great! Let's gather requirements."
assert body["draft_ready"] is False
@pytest.mark.asyncio
async def test_send_message_marks_draft_ready(prompter_client: dict) -> None:
"""A ready roboco-meta control block flips draft_ready and session status."""
client = prompter_client["client"]
session_resp = await client.post("/api/prompter/sessions", json={}, headers=_HDR)
session_id = session_resp.json()["id"]
mock_response = (
"Understood — I have what I need.\n\n"
'```roboco-meta\n{"covered": ["objective", "scope", "surface", '
'"acceptance"], "ready": true, "scale": "single"}\n```'
)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=mock_response,
):
response = await client.post(
f"/api/prompter/sessions/{session_id}/messages",
json={"content": "Add a login page with MFA support"},
headers=_HDR,
)
assert response.status_code == HTTPStatus.OK
body = response.json()
assert len(body["messages"]) == _SINGLE_TURN_MSGS
assert body["draft_ready"] is True
assert body["scale"] == "single"
# The control block must not leak into the persisted assistant message.
assert "roboco-meta" not in body["messages"][-1]["content"]
@pytest.mark.asyncio
async def test_send_message_not_found(prompter_client: dict) -> None:
"""POST /sessions/{id}/messages with unknown session → 404."""
client = prompter_client["client"]
response = await client.post(
f"/api/prompter/sessions/{uuid4()}/messages",
json={"content": "Hello"},
headers=_HDR,
)
assert response.status_code == HTTPStatus.NOT_FOUND
@pytest.mark.asyncio
async def test_get_draft_generates_from_conversation(prompter_client: dict) -> None:
"""GET /sessions/{id}/draft generates a draft via LLM."""
client = prompter_client["client"]
session_resp = await client.post("/api/prompter/sessions", json={}, headers=_HDR)
session_id = session_resp.json()["id"]
draft_json = {
"title": "Add login page",
"description": "Implement a secure login page with email and password",
"acceptance_criteria": [
"User can enter email and password",
"Invalid credentials show error message",
],
"team": "frontend",
"task_type": "code",
"nature": "technical",
"estimated_complexity": "medium",
"priority": 2,
}
chat_response = "Tell me more about the requirements."
draft_response = json.dumps(draft_json)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=chat_response,
):
await client.post(
f"/api/prompter/sessions/{session_id}/messages",
json={"content": "I need a login page"},
headers=_HDR,
)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=draft_response,
):
response = await client.get(
f"/api/prompter/sessions/{session_id}/draft",
headers=_HDR,
)
assert response.status_code == HTTPStatus.OK
body = response.json()
assert body["draft"]["title"] == "Add login page"
assert body["draft"]["source"] == "prompter"
assert body["confirmed_at"] is None
assert body["draft"]["confirmed_by_human"] is False
assert body["session_id"] == session_id
@pytest.mark.asyncio
async def test_get_draft_cached(prompter_client: dict) -> None:
"""GET /sessions/{id}/draft returns the cached draft on subsequent calls."""
client = prompter_client["client"]
session_resp = await client.post("/api/prompter/sessions", json={}, headers=_HDR)
session_id = session_resp.json()["id"]
draft_json = {
"title": "Add login page",
"description": "Implement a secure login page with email and password",
"acceptance_criteria": ["User can enter credentials"],
"team": "frontend",
"task_type": "code",
"nature": "technical",
"estimated_complexity": "medium",
"priority": 2,
}
chat_response = "Got it."
draft_response = json.dumps(draft_json)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=chat_response,
):
await client.post(
f"/api/prompter/sessions/{session_id}/messages",
json={"content": "I need a login page"},
headers=_HDR,
)
call_count = 0
async def _mock_create(**_kwargs: Any) -> str:
nonlocal call_count
call_count += 1
return draft_response
with patch(
"roboco.services.prompter.PrompterService._create_message",
side_effect=_mock_create,
):
await client.get(f"/api/prompter/sessions/{session_id}/draft", headers=_HDR)
second_response = await client.get(
f"/api/prompter/sessions/{session_id}/draft", headers=_HDR
)
assert second_response.status_code == HTTPStatus.OK
# LLM should only be called once (draft is cached)
assert call_count == 1
@pytest.mark.asyncio
async def test_get_draft_empty_session_returns_400(prompter_client: dict) -> None:
"""GET /sessions/{id}/draft with no messages → 400."""
client = prompter_client["client"]
session_resp = await client.post("/api/prompter/sessions", json={}, headers=_HDR)
session_id = session_resp.json()["id"]
response = await client.get(
f"/api/prompter/sessions/{session_id}/draft",
headers=_HDR,
)
assert response.status_code == HTTPStatus.BAD_REQUEST
@pytest.mark.asyncio
async def test_confirm_draft_creates_task(
prompter_client: dict, project_fixture: ProjectTable
) -> None:
"""POST /sessions/{id}/confirm validates draft and creates a real task."""
client = prompter_client["client"]
project_id = str(project_fixture.id)
session_resp = await client.post("/api/prompter/sessions", json={}, headers=_HDR)
session_id = session_resp.json()["id"]
draft_json = {
"title": "Add login page",
"description": "Implement a secure login page with email and password",
"acceptance_criteria": ["User can enter credentials"],
"team": "frontend",
"task_type": "code",
"nature": "technical",
"estimated_complexity": "medium",
"priority": 2,
}
chat_response = "Got it."
draft_response = json.dumps(draft_json)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=chat_response,
):
await client.post(
f"/api/prompter/sessions/{session_id}/messages",
json={"content": "I need a login page"},
headers=_HDR,
)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=draft_response,
):
await client.get(f"/api/prompter/sessions/{session_id}/draft", headers=_HDR)
confirm_response = await client.post(
f"/api/prompter/sessions/{session_id}/confirm",
json={"project_id": project_id},
headers=_HDR,
)
assert confirm_response.status_code == HTTPStatus.CREATED
body = confirm_response.json()
assert "task_id" in body
assert body["task_id"] is not None
@pytest.mark.asyncio
async def test_confirm_draft_requires_project_or_product(
prompter_client: dict,
) -> None:
"""POST /sessions/{id}/confirm without project_id/product_id → 400."""
client = prompter_client["client"]
session_resp = await client.post("/api/prompter/sessions", json={}, headers=_HDR)
session_id = session_resp.json()["id"]
draft_json = {
"title": "Add login page",
"description": "Implement a secure login page with email and password",
"acceptance_criteria": ["User can enter credentials"],
"team": "frontend",
"task_type": "code",
"nature": "technical",
"estimated_complexity": "medium",
"priority": 2,
}
chat_response = "Got it."
draft_response = json.dumps(draft_json)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=chat_response,
):
await client.post(
f"/api/prompter/sessions/{session_id}/messages",
json={"content": "I need a login page"},
headers=_HDR,
)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=draft_response,
):
await client.get(f"/api/prompter/sessions/{session_id}/draft", headers=_HDR)
confirm_response = await client.post(
f"/api/prompter/sessions/{session_id}/confirm",
json={},
headers=_HDR,
)
assert confirm_response.status_code == HTTPStatus.BAD_REQUEST
# =============================================================================
# Full happy path integration test
# =============================================================================
@pytest.mark.asyncio
async def test_full_happy_path(
prompter_client: dict, project_fixture: ProjectTable
) -> None:
"""Full happy path: create session → send messages → get draft → confirm task."""
client = prompter_client["client"]
project_id = str(project_fixture.id)
# Step 1: Create session
step1 = await client.post("/api/prompter/sessions", json={}, headers=_HDR)
assert step1.status_code == HTTPStatus.CREATED
session_id = step1.json()["id"]
# Step 2: Send messages
chat_mock = "Please describe the acceptance criteria for this feature."
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=chat_mock,
):
step2a = await client.post(
f"/api/prompter/sessions/{session_id}/messages",
json={"content": "I need a dark mode toggle for the UI"},
headers=_HDR,
)
assert step2a.status_code == HTTPStatus.OK
chat_mock2 = "I have enough information to draft a task now."
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=chat_mock2,
):
step2b = await client.post(
f"/api/prompter/sessions/{session_id}/messages",
json={"content": "Preference is persisted across sessions"},
headers=_HDR,
)
assert step2b.status_code == HTTPStatus.OK
messages = step2b.json()["messages"]
assert len(messages) == _DOUBLE_TURN_MSGS
# Step 3: Get draft
draft_json = {
"title": "Add dark mode toggle",
"description": "Implement a dark mode toggle so users can switch themes",
"acceptance_criteria": [
"User can toggle light/dark mode",
"Preference is persisted across sessions",
],
"team": "frontend",
"task_type": "code",
"nature": "technical",
"estimated_complexity": "low",
"priority": 2,
}
draft_mock = json.dumps(draft_json)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=draft_mock,
):
step3 = await client.get(
f"/api/prompter/sessions/{session_id}/draft", headers=_HDR
)
assert step3.status_code == HTTPStatus.OK
draft_body = step3.json()
assert draft_body["draft"]["title"] == "Add dark mode toggle"
# Step 4: Confirm draft → creates task
step4 = await client.post(
f"/api/prompter/sessions/{session_id}/confirm",
json={"project_id": project_id},
headers=_HDR,
)
assert step4.status_code == HTTPStatus.CREATED
task_body = step4.json()
assert "task_id" in task_body
assert task_body["task_id"] is not None
# =============================================================================
# Legacy stateless endpoint tests (backward compatibility)
# =============================================================================
@pytest.mark.asyncio
async def test_prompter_chat_success(prompter_client: dict) -> None:
client = prompter_client["client"]
mock_response = "Great! Let's gather requirements."
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=mock_response,
):
response = await client.post(
"/api/prompter/chat",
json={
"messages": [{"role": "user", "content": "I need a new feature"}],
},
headers=_HDR,
)
assert response.status_code == HTTPStatus.OK
body = response.json()
assert body["message"] == "Great! Let's gather requirements."
assert body["draft_ready"] is False
@pytest.mark.asyncio
async def test_prompter_chat_draft_ready(prompter_client: dict) -> None:
client = prompter_client["client"]
mock_response = (
"Understood.\n\n"
'```roboco-meta\n{"covered": ["objective", "scope", "surface", '
'"acceptance"], "ready": true, "scale": "single"}\n```'
)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=mock_response,
):
response = await client.post(
"/api/prompter/chat",
json={
"messages": [
{"role": "user", "content": "I need a new feature"},
{"role": "assistant", "content": "Tell me more"},
{"role": "user", "content": "Add a login page"},
],
},
headers=_HDR,
)
assert response.status_code == HTTPStatus.OK
body = response.json()
assert body["draft_ready"] is True
@pytest.mark.asyncio
async def test_prompter_chat_llm_failure(prompter_client: dict) -> None:
client = prompter_client["client"]
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
side_effect=Exception("Anthropic API unavailable"),
):
response = await client.post(
"/api/prompter/chat",
json={
"messages": [{"role": "user", "content": "Hello"}],
},
headers=_HDR,
)
assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
body = response.json()
assert "LLM chat failed" in body["detail"]["message"]
@pytest.mark.asyncio
async def test_prompter_draft_success(prompter_client: dict) -> None:
client = prompter_client["client"]
draft_json = {
"title": "Add login page",
"description": "Implement a secure login page with email and password",
"acceptance_criteria": [
"User can enter email and password",
"Invalid credentials show error message",
],
"team": "frontend",
"task_type": "code",
"nature": "technical",
"estimated_complexity": "medium",
"priority": 2,
}
mock_response = json.dumps(draft_json)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=mock_response,
):
response = await client.post(
"/api/prompter/draft",
json={
"messages": [
{"role": "user", "content": "I need a login page"},
],
},
headers=_HDR,
)
assert response.status_code == HTTPStatus.OK
body = response.json()
assert body["draft"]["title"] == "Add login page"
assert body["draft"]["source"] == "prompter"
assert body["draft"]["confirmed_by_human"] is False
assert "reasoning" in body
@pytest.mark.asyncio
async def test_prompter_draft_invalid_json_from_llm(prompter_client: dict) -> None:
client = prompter_client["client"]
mock_response = "not valid json"
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=mock_response,
):
response = await client.post(
"/api/prompter/draft",
json={
"messages": [{"role": "user", "content": "Hello"}],
},
headers=_HDR,
)
assert response.status_code == HTTPStatus.BAD_REQUEST
body = response.json()
assert "Draft response was not valid JSON" in body["detail"]["message"]
@pytest.mark.asyncio
async def test_prompter_draft_schema_mismatch(prompter_client: dict) -> None:
client = prompter_client["client"]
bad_draft = {
"title": "x",
"description": "too short",
}
mock_response = json.dumps(bad_draft)
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
return_value=mock_response,
):
response = await client.post(
"/api/prompter/draft",
json={
"messages": [{"role": "user", "content": "Hello"}],
},
headers=_HDR,
)
assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
body = response.json()
assert "draft_schema_error" in body["detail"]["error"]
@pytest.mark.asyncio
async def test_prompter_draft_llm_failure(prompter_client: dict) -> None:
client = prompter_client["client"]
with patch(
"roboco.services.prompter.PrompterService._create_message",
new_callable=AsyncMock,
side_effect=Exception("Anthropic API unavailable"),
):
response = await client.post(
"/api/prompter/draft",
json={
"messages": [{"role": "user", "content": "Hello"}],
},
headers=_HDR,
)
assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
body = response.json()
assert "LLM draft generation failed" in body["detail"]["message"]
@pytest.mark.asyncio
async def test_prompter_chat_empty_messages(prompter_client: dict) -> None:
client = prompter_client["client"]
response = await client.post(
"/api/prompter/chat",
json={"messages": []},
headers=_HDR,
)
assert response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
@pytest.mark.asyncio
async def test_prompter_chat_invalid_role(prompter_client: dict) -> None:
client = prompter_client["client"]
response = await client.post(
"/api/prompter/chat",
json={"messages": [{"role": "invalid", "content": "hi"}]},
headers=_HDR,
)
assert response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
@pytest.mark.asyncio
async def test_prompter_draft_empty_messages(prompter_client: dict) -> None:
client = prompter_client["client"]
response = await client.post(
"/api/prompter/draft",
json={"messages": []},
headers=_HDR,
)
assert response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
tests/unit/agents/test_tool_load_directive_layer.py
+18
View File
@@ -95,6 +95,24 @@ def test_pm_blocks_exclude_edit_and_write() -> None:
)
def test_no_role_is_granted_the_task_subagent_tool() -> None:
"""Task (sub-agent dispatch) is dropped from the built-in tool grant.
No role prompt or workflow uses Task and there are no custom sub-agent
definitions, so a Task call only spawns a context-blind generic sub-agent
that burns budget. The tools-ready line must not advertise it for any role.
"""
for role in (
AgentRole.DEVELOPER,
AgentRole.DOCUMENTER,
AgentRole.QA,
AgentRole.MAIN_PM,
AgentRole.CELL_PM,
):
names = _tool_names(_composed_prompt_for(role, Team.BACKEND))
assert "Task" not in names, f"{role.value} must not list Task: {names}"
def test_block_is_first_layer_before_lifecycle() -> None:
"""Tools-ready block precedes the lifecycle and base layers."""
prompt = _composed_prompt_for(AgentRole.DEVELOPER, Team.BACKEND)
tests/unit/api/test_schemas_prompter.py
-272
View File
@@ -1,272 +0,0 @@
"""Unit tests for Prompter API schemas.
Covers schema validation for both the session-based and legacy schemas.
"""
from __future__ import annotations
from uuid import uuid4
import pytest
from pydantic import ValidationError as PydanticValidationError
from roboco.api.schemas.prompter import (
CellWork,
ChatMessage,
PrompterChatRequest,
PrompterDraftTask,
PrompterMessageRequest,
PrompterTurnResponse,
TaskConfirmRequest,
)
# =============================================================================
# ChatMessage
# =============================================================================
def test_chat_message_valid_roles() -> None:
for role in ("user", "assistant", "system"):
msg = ChatMessage(role=role, content="Hello")
assert msg.role == role
def test_chat_message_invalid_role() -> None:
with pytest.raises(PydanticValidationError) as exc_info:
ChatMessage(role="admin", content="Hello")
assert "role must be one of" in str(exc_info.value)
def test_chat_message_empty_content() -> None:
with pytest.raises(PydanticValidationError):
ChatMessage(role="user", content="")
# =============================================================================
# PrompterMessageRequest
# =============================================================================
def test_message_request_valid() -> None:
req = PrompterMessageRequest(content="I need a feature")
assert req.content == "I need a feature"
assert req.context == {}
def test_message_request_empty_content() -> None:
with pytest.raises(PydanticValidationError):
PrompterMessageRequest(content="")
def test_message_request_with_context() -> None:
req = PrompterMessageRequest(content="Hello", context={"key": "value"})
assert req.context["key"] == "value"
# =============================================================================
# TaskConfirmRequest
# =============================================================================
def test_task_confirm_request_all_optional() -> None:
req = TaskConfirmRequest()
assert req.project_id is None
assert req.product_id is None
assert req.assigned_to is None
assert req.overrides == {}
def test_task_confirm_request_with_project() -> None:
pid = uuid4()
req = TaskConfirmRequest(project_id=pid)
assert req.project_id == pid
# =============================================================================
# PrompterDraftTask
# =============================================================================
def test_draft_task_valid() -> None:
draft = PrompterDraftTask(
title="Add login page",
description="Implement a secure login page with email and password",
acceptance_criteria=["User can log in"],
team="frontend",
task_type="code",
nature="technical",
estimated_complexity="medium",
)
assert draft.title == "Add login page"
assert draft.source == "prompter"
assert draft.confirmed_by_human is False
def test_draft_task_title_too_long() -> None:
with pytest.raises(PydanticValidationError):
PrompterDraftTask(
title="x" * 201,
description="Implement a secure login page with email and password",
acceptance_criteria=["User can log in"],
team="frontend",
task_type="code",
nature="technical",
estimated_complexity="medium",
)
def test_draft_task_description_too_short() -> None:
with pytest.raises(PydanticValidationError):
PrompterDraftTask(
title="Add login page",
description="short", # <20 chars
acceptance_criteria=["User can log in"],
team="frontend",
task_type="code",
nature="technical",
estimated_complexity="medium",
)
def test_draft_task_empty_acceptance_criteria() -> None:
with pytest.raises(PydanticValidationError):
PrompterDraftTask(
title="Add login page",
description="Implement a secure login page with email and password",
acceptance_criteria=[],
team="frontend",
task_type="code",
nature="technical",
estimated_complexity="medium",
)
def test_draft_task_invalid_team() -> None:
with pytest.raises(PydanticValidationError):
PrompterDraftTask(
title="Add login page",
description="Implement a secure login page with email and password",
acceptance_criteria=["User can log in"],
team="infra", # invalid
task_type="code",
nature="technical",
estimated_complexity="medium",
)
def test_draft_task_priority_bounds() -> None:
# Valid bounds
for p in (0, 1, 2, 3):
d = PrompterDraftTask(
title="Add login page",
description="Implement a secure login page with email and password",
acceptance_criteria=["User can log in"],
team="frontend",
task_type="code",
nature="technical",
estimated_complexity="medium",
priority=p,
)
assert d.priority == p
# Out of bounds
with pytest.raises(PydanticValidationError):
PrompterDraftTask(
title="Add login page",
description="Implement a secure login page with email and password",
acceptance_criteria=["User can log in"],
team="frontend",
task_type="code",
nature="technical",
estimated_complexity="medium",
priority=4,
)
# =============================================================================
# Structured spec fields
# =============================================================================
def test_cell_work_valid() -> None:
cw = CellWork(team="backend", summary="Build the endpoint", items=["Route", "Test"])
assert cw.team.value == "backend"
assert cw.items == ["Route", "Test"]
def test_cell_work_requires_summary() -> None:
with pytest.raises(PydanticValidationError):
CellWork(team="backend", summary="")
def test_draft_task_structured_fields_default_empty() -> None:
draft = PrompterDraftTask(
title="Add login page",
description="Implement a secure login page with email and password",
acceptance_criteria=["User can log in"],
team="frontend",
task_type="code",
nature="technical",
estimated_complexity="medium",
)
assert draft.objective is None
assert draft.what_this_builds == []
assert draft.the_work == []
assert draft.notes == []
def test_draft_task_with_structured_fields() -> None:
draft = PrompterDraftTask(
title="Ship the Prompter",
description="A board-led feature spanning three cells, fully wired.",
acceptance_criteria=["It works end to end"],
team="backend",
task_type="code",
nature="technical",
estimated_complexity="high",
objective="Let humans chat a task into existence.",
what_this_builds=["A /prompter page", "A chat endpoint"],
the_work=[
CellWork(team="backend", summary="Chat endpoint", items=["Route"]),
CellWork(team="frontend", summary="Chat UI", items=["Page"]),
],
notes=["Reuse the LLM service"],
)
assert [w.team.value for w in draft.the_work] == ["backend", "frontend"]
def test_confirm_request_carries_edited_draft() -> None:
draft = PrompterDraftTask(
title="Edited title",
description="An edited description that clears the minimum length.",
acceptance_criteria=["Done"],
team="frontend",
task_type="code",
nature="technical",
estimated_complexity="low",
)
req = TaskConfirmRequest(project_id=uuid4(), draft=draft)
assert req.draft is not None
assert req.draft.title == "Edited title"
def test_turn_response_shape() -> None:
resp = PrompterTurnResponse(messages=[], draft_ready=True, scale="multi")
assert resp.draft_ready is True
assert resp.scale == "multi"
assert resp.messages == []
# =============================================================================
# PrompterChatRequest (legacy)
# =============================================================================
def test_chat_request_requires_messages() -> None:
with pytest.raises(PydanticValidationError):
PrompterChatRequest(messages=[])
def test_chat_request_valid() -> None:
req = PrompterChatRequest(messages=[ChatMessage(role="user", content="Hello")])
assert len(req.messages) == 1
assert req.context == {}
tests/unit/api/test_websocket_bridge.py
+1 -29
View File
@@ -297,33 +297,6 @@ async def test_handle_rate_limit_ignores_unrelated_event() -> None:
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_handle_usage_update_broadcasts_to_system() -> None:
"""USAGE_UPDATE event → broadcast_system tagged USAGE_UPDATE + data fields."""
expected_input = 100
expected_output = 50
event = _evt(
EventType.USAGE_UPDATE,
{
"agent_id": "be-dev-1",
"task_id": "task-abc",
"input_tokens": expected_input,
"output_tokens": expected_output,
"model": "claude-sonnet-4-6",
"timestamp": "2026-06-11T00:00:00+00:00",
},
)
with patch("roboco.api.websocket_bridge.manager") as mgr:
mgr.broadcast_system = AsyncMock()
await _handle_usage_event(event)
mgr.broadcast_system.assert_awaited_once()
msg = mgr.broadcast_system.await_args.args[0]
assert msg["type"] == "USAGE_UPDATE"
assert msg["agent_id"] == "be-dev-1"
assert msg["input_tokens"] == expected_input
assert msg["output_tokens"] == expected_output
@pytest.mark.asyncio
async def test_handle_usage_snapshot_broadcasts_to_system() -> None:
"""USAGE_SNAPSHOT event → broadcast_system tagged USAGE_SNAPSHOT + aggregate."""
@@ -376,7 +349,7 @@ def test_register_websocket_bridge_handlers_subscribes_all_event_types() -> None
with patch("roboco.api.websocket_bridge.get_event_bus", return_value=fake):
register_websocket_bridge_handlers()
types = [t for t, _ in fake.subscribed]
# All 14 expected event types appear at least once.
# All expected event types appear at least once.
assert EventType.NOTIFICATION_SENT in types
assert EventType.NOTIFICATION_ACKED in types
assert EventType.SESSION_CREATED in types
@@ -390,7 +363,6 @@ def test_register_websocket_bridge_handlers_subscribes_all_event_types() -> None
assert EventType.RATE_LIMIT_HIT in types
assert EventType.RATE_LIMIT_LIFTED in types
# Usage events forwarded to /ws/system.
assert EventType.USAGE_UPDATE in types
assert EventType.USAGE_SNAPSHOT in types
tests/unit/runtime/test_briefing_tool_load_block.py
+20
View File
@@ -68,6 +68,26 @@ def test_pm_blocks_exclude_edit_and_write() -> None:
assert "Write" not in names, f"{role} must not list Write"
def test_no_role_block_lists_the_task_subagent_tool() -> None:
"""Task (sub-agent dispatch) is dropped from the briefing tool grant.
No role uses Task and there are no custom sub-agent definitions, so it only
spawns a context-blind generic sub-agent that burns budget.
"""
for role in (
"developer",
"documenter",
"qa",
"main_pm",
"cell_pm",
"product_owner",
"head_marketing",
"auditor",
):
names = _tool_names(_orch()._build_tool_load_block(role))
assert "Task" not in names, f"{role} must not list Task: {names}"
def test_unknown_role_returns_empty() -> None:
assert _orch()._build_tool_load_block("nonexistent") == ""
tests/unit/services/test_prompter.py
+6 -251
View File
@@ -1,14 +1,14 @@
"""Unit tests for PrompterService.
Tests the service layer logic with mocked LLM calls. Uses an in-memory
async session (via conftest fixtures) for DB-backed tests.
Covers the live-intake draft → task flow (``create_task_from_draft`` /
``confirm_live_draft`` + the enum/priority/team coercion) and the pure
draft/description helpers. DB-backed tests use an in-memory async session via
conftest fixtures.
"""
from __future__ import annotations
import json
from typing import Any
from unittest.mock import AsyncMock, patch
from uuid import UUID, uuid4
import pytest
@@ -28,12 +28,9 @@ from roboco.models.base import (
Team,
)
from roboco.seeds.initial_data import AGENT_UUIDS
from roboco.services.base import NotFoundError, ServiceError, ValidationError
from roboco.services.base import ServiceError
from roboco.services.prompter import (
PrompterService,
_build_chat_prompt,
_build_draft_prompt,
_build_reasoning,
compose_description,
derive_scale,
get_prompter_service,
@@ -209,41 +206,6 @@ def test_coerce_draft_enums_keeps_valid_and_derives_missing_team() -> None:
assert complexity is Complexity.MEDIUM
def test_build_chat_prompt_basic() -> None:
messages = [
{"role": "user", "content": "I need a feature"},
{"role": "assistant", "content": "Tell me more"},
]
prompt = _build_chat_prompt(messages, None)
assert "user: I need a feature" in prompt
assert "assistant: Tell me more" in prompt
assert "Continue the conversation" in prompt
def test_build_chat_prompt_with_context() -> None:
messages = [{"role": "user", "content": "hello"}]
prompt = _build_chat_prompt(messages, {"team": "backend"})
assert "Context:" in prompt
assert "team: backend" in prompt
def test_build_draft_prompt() -> None:
messages = [{"role": "user", "content": "I need a login page"}]
prompt = _build_draft_prompt(messages, None)
assert "valid JSON" in prompt
assert "user: I need a login page" in prompt
def test_build_reasoning() -> None:
messages = [{"role": "user", "content": "Hello"}] * 3
draft = {"title": "My Task", "team": "backend", "estimated_complexity": "medium"}
reasoning = _build_reasoning(messages, draft)
assert "My Task" in reasoning
assert "backend" in reasoning
assert "medium" in reasoning
assert "3 messages" in reasoning
# =============================================================================
# Factory
# =============================================================================
@@ -262,178 +224,10 @@ def test_get_prompter_service_raises_without_db_for_session_methods() -> None:
# =============================================================================
# Stateless chat / draft (with mocked LLM)
# DB-backed: assignee routing + confirm_live_draft
# =============================================================================
@pytest.mark.asyncio
async def test_chat_success_with_mock_llm() -> None:
service = get_prompter_service()
with patch.object(
service,
"_create_message",
new_callable=AsyncMock,
return_value="Great, let's continue!",
):
result = await service.chat(
messages=[{"role": "user", "content": "I need a feature"}]
)
assert result["message"] == "Great, let's continue!"
assert result["draft_ready"] is False
@pytest.mark.asyncio
async def test_chat_draft_ready_signal() -> None:
service = get_prompter_service()
reply = (
"Got it — I have what I need.\n\n"
'```roboco-meta\n{"covered": ["objective", "scope", "surface", '
'"acceptance"], "ready": true, "scale": "single"}\n```'
)
with patch.object(
service,
"_create_message",
new_callable=AsyncMock,
return_value=reply,
):
result = await service.chat(
messages=[{"role": "user", "content": "I need a feature"}]
)
assert result["draft_ready"] is True
assert result["scale"] == "single"
# The control block is stripped from the user-visible reply.
assert "roboco-meta" not in result["message"]
assert result["message"] == "Got it — I have what I need."
@pytest.mark.asyncio
async def test_chat_raises_on_empty_response() -> None:
service = get_prompter_service()
with (
patch.object(
service, "_create_message", new_callable=AsyncMock, return_value=""
),
pytest.raises(ServiceError, match="LLM returned empty content"),
):
await service.chat(messages=[{"role": "user", "content": "Hello"}])
@pytest.mark.asyncio
async def test_chat_raises_on_llm_error() -> None:
service = get_prompter_service()
with (
patch.object(
service,
"_create_message",
new_callable=AsyncMock,
side_effect=Exception("API unavailable"),
),
pytest.raises(ServiceError, match="LLM chat failed"),
):
await service.chat(messages=[{"role": "user", "content": "Hello"}])
@pytest.mark.asyncio
async def test_draft_success_with_mock_llm() -> None:
service = get_prompter_service()
draft_data = {
"title": "Add login",
"description": "Implement login functionality with JWT tokens",
"acceptance_criteria": ["User can log in"],
"team": "backend",
"task_type": "code",
"nature": "technical",
"estimated_complexity": "medium",
}
with patch.object(
service,
"_create_message",
new_callable=AsyncMock,
return_value=json.dumps(draft_data),
):
result = await service.draft(
messages=[{"role": "user", "content": "I need a login feature"}]
)
assert result["draft"]["title"] == "Add login"
assert result["draft"]["source"] == "prompter"
assert result["draft"]["confirmed_by_human"] is False
assert "reasoning" in result
@pytest.mark.asyncio
async def test_draft_raises_on_invalid_json() -> None:
service = get_prompter_service()
with (
patch.object(
service,
"_create_message",
new_callable=AsyncMock,
return_value="Not JSON at all",
),
pytest.raises(ValidationError, match="not valid JSON"),
):
await service.draft(messages=[{"role": "user", "content": "Hello"}])
@pytest.mark.asyncio
async def test_draft_raises_on_llm_error() -> None:
service = get_prompter_service()
with (
patch.object(
service,
"_create_message",
new_callable=AsyncMock,
side_effect=Exception("API unavailable"),
),
pytest.raises(ServiceError, match="LLM draft generation failed"),
):
await service.draft(messages=[{"role": "user", "content": "Hello"}])
# =============================================================================
# Session-based: create_session (DB-backed via conftest)
# =============================================================================
@pytest.mark.asyncio
async def test_create_session_db(db_session: Any) -> None:
"""create_session persists a PrompterSessionTable row."""
service = get_prompter_service(db=db_session)
agent_id = uuid4()
agent = AgentTable(
id=agent_id,
name="TestAgent",
slug=f"test-{uuid4().hex[:8]}",
role=AgentRole.DEVELOPER,
team=None,
status=AgentStatus.ACTIVE,
model_config={},
system_prompt="dev",
capabilities=[],
permissions={},
metrics={},
)
db_session.add(agent)
await db_session.flush()
session = await service.create_session(agent_id=agent_id)
assert session.id is not None
assert session.status == "active"
assert session.agent_id == agent_id
@pytest.mark.asyncio
async def test_assignee_is_board_distinguishes_roles(db_session: Any) -> None:
"""Drives product team routing: a board reviewer keeps the root on the board.
@@ -472,45 +266,6 @@ async def test_assignee_is_board_distinguishes_roles(db_session: Any) -> None:
assert await service._assignee_is_board(uuid4()) is False
@pytest.mark.asyncio
async def test_get_session_not_found(db_session: Any) -> None:
"""_get_session raises NotFoundError for unknown session ID."""
service = get_prompter_service(db=db_session)
with pytest.raises(NotFoundError):
await service._get_session(uuid4(), uuid4())
@pytest.mark.asyncio
async def test_get_draft_empty_session_raises(db_session: Any) -> None:
"""get_or_generate_draft raises ValidationError if no messages exist."""
service = get_prompter_service(db=db_session)
agent_id = uuid4()
agent = AgentTable(
id=agent_id,
name="TestAgent",
slug=f"test-{uuid4().hex[:8]}",
role=AgentRole.DEVELOPER,
team=None,
status=AgentStatus.ACTIVE,
model_config={},
system_prompt="dev",
capabilities=[],
permissions={},
metrics={},
)
db_session.add(agent)
await db_session.flush()
session = await service.create_session(agent_id=agent_id)
with pytest.raises(ValidationError, match="empty conversation"):
await service.get_or_generate_draft(
session_id=UUID(str(session.id)),
agent_id=agent_id,
)
async def _seed_project_and_ceo(db_session: Any) -> tuple[UUID, UUID]:
"""Seed a system agent + project + CEO; return (project_id, ceo_id).
tests/unit/services/test_usage_events.py
+4 -184
View File
@@ -1,196 +1,16 @@
"""Unit tests for roboco.services.usage_events.
Covers the _UsageThrottle class and the publish_usage_update /
publish_usage_snapshot helpers. No real Redis or event bus is needed —
we use AsyncMock to assert that bus.publish is called with the right
Covers the publish_usage_snapshot helper. No real Redis or event bus is
needed — we use AsyncMock to assert that bus.publish is called with the right
payload and type.
The throttle suppression test is the acceptance-criterion gate:
"Server-side throttle prevents more than 1 USAGE_UPDATE publish per
agent per 5-second window."
"""
from __future__ import annotations
from datetime import UTC, datetime
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, MagicMock
import pytest
from roboco.services.usage_events import (
UsageSnapshot,
UsageUpdate,
_UsageThrottle,
publish_usage_snapshot,
publish_usage_update,
)
# ---------------------------------------------------------------------------
# _UsageThrottle
# ---------------------------------------------------------------------------
def test_throttle_allows_first_publish() -> None:
"""A fresh agent has no prior timestamp — first publish is always allowed."""
th = _UsageThrottle(window=5.0)
assert th.should_publish("be-dev-1") is True
def test_throttle_suppresses_second_publish_within_window() -> None:
"""Second call within the 5-second window returns False (suppressed)."""
th = _UsageThrottle(window=5.0)
with patch("roboco.services.usage_events.time") as mock_time:
mock_time.monotonic.return_value = 100.0
assert th.should_publish("be-dev-1") is True # first → allowed
mock_time.monotonic.return_value = 104.9 # 4.9 s later — still inside window
assert th.should_publish("be-dev-1") is False # suppressed
def test_throttle_allows_publish_after_window_expires() -> None:
"""After the full window elapses, the next publish is allowed again."""
th = _UsageThrottle(window=5.0)
with patch("roboco.services.usage_events.time") as mock_time:
mock_time.monotonic.return_value = 100.0
assert th.should_publish("be-dev-1") is True # first
mock_time.monotonic.return_value = 105.0 # exactly 5 s later
assert th.should_publish("be-dev-1") is True # window elapsed → allowed
def test_throttle_tracks_agents_independently() -> None:
"""Different agents have independent throttle windows."""
th = _UsageThrottle(window=5.0)
with patch("roboco.services.usage_events.time") as mock_time:
mock_time.monotonic.return_value = 100.0
assert th.should_publish("be-dev-1") is True
# be-dev-2 has never published, so it is always allowed.
assert th.should_publish("be-dev-2") is True
mock_time.monotonic.return_value = 101.0
# be-dev-1 is suppressed; be-dev-2 is also now suppressed.
assert th.should_publish("be-dev-1") is False
assert th.should_publish("be-dev-2") is False
def test_throttle_records_timestamp_on_allow() -> None:
"""should_publish records the current time when it returns True."""
th = _UsageThrottle(window=5.0)
recorded_at = 200.0
with patch("roboco.services.usage_events.time") as mock_time:
mock_time.monotonic.return_value = recorded_at
th.should_publish("be-dev-1")
assert th._last["be-dev-1"] == recorded_at
# ---------------------------------------------------------------------------
# publish_usage_update
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_publish_usage_update_calls_bus_publish() -> None:
"""First call in a window publishes the event and returns True."""
bus = MagicMock()
bus.publish = AsyncMock()
th = _UsageThrottle(window=5.0)
expected_input = 100
expected_output = 50
with patch("roboco.services.usage_events._throttle", th):
result = await publish_usage_update(
bus,
UsageUpdate(
agent_id="be-dev-1",
task_id="task-abc",
input_tokens=expected_input,
output_tokens=expected_output,
model="claude-sonnet-4-6",
),
)
assert result is True
bus.publish.assert_awaited_once()
event = bus.publish.await_args.args[0]
assert event.type.value == "usage.update"
assert event.data["agent_id"] == "be-dev-1"
assert event.data["task_id"] == "task-abc"
assert event.data["input_tokens"] == expected_input
assert event.data["output_tokens"] == expected_output
assert event.data["model"] == "claude-sonnet-4-6"
assert "timestamp" in event.data
@pytest.mark.asyncio
async def test_publish_usage_update_throttle_suppresses_second_call() -> None:
"""Second publish within the throttle window is suppressed (returns False)."""
bus = MagicMock()
bus.publish = AsyncMock()
th = _UsageThrottle(window=5.0)
with (
patch("roboco.services.usage_events._throttle", th),
patch("roboco.services.usage_events.time") as mock_time,
):
mock_time.monotonic.return_value = 100.0
first = await publish_usage_update(
bus,
UsageUpdate(
agent_id="be-dev-1",
task_id=None,
input_tokens=10,
output_tokens=5,
model="sonnet",
),
)
mock_time.monotonic.return_value = 102.0 # 2 s later — still suppressed
second = await publish_usage_update(
bus,
UsageUpdate(
agent_id="be-dev-1",
task_id=None,
input_tokens=20,
output_tokens=10,
model="sonnet",
),
)
assert first is True
assert second is False
# bus.publish should only have been called once.
assert bus.publish.await_count == 1
@pytest.mark.asyncio
async def test_publish_usage_update_custom_timestamp() -> None:
"""Custom timestamp is passed through to the event data."""
bus = MagicMock()
bus.publish = AsyncMock()
ts = datetime(2026, 6, 11, 12, 0, 0, tzinfo=UTC)
# Use a fresh throttle so the first publish goes through.
th = _UsageThrottle(window=5.0)
with patch("roboco.services.usage_events._throttle", th):
await publish_usage_update(
bus,
UsageUpdate(
agent_id="be-dev-1",
task_id=None,
input_tokens=0,
output_tokens=0,
model="sonnet",
timestamp=ts,
),
)
event = bus.publish.await_args.args[0]
assert event.data["timestamp"] == ts.isoformat()
from roboco.services.usage_events import UsageSnapshot, publish_usage_snapshot
# ---------------------------------------------------------------------------
# publish_usage_snapshot