diff --git a/duck_core/api.py b/duck_core/api.py index 49f92eb..4c78b46 100644 --- a/duck_core/api.py +++ b/duck_core/api.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import time from pathlib import Path from typing import Any @@ -13,7 +14,7 @@ from pydantic import BaseModel from duck_core.approvals.service import ApprovalService from duck_core.config import get_settings -from duck_core.context_builder import ContextBuilder +from duck_core.context_builder import ContextBuilder, estimate_tokens from duck_core.conversations.store import ConversationStore from duck_core.events.store import EventStore from duck_core.experience.recorder import ExperienceRecorder @@ -28,6 +29,32 @@ from duck_core.tasks.store import TaskStore logger = logging.getLogger(__name__) +class GenerationStats: + def __init__(self) -> None: + self.started_at = time.perf_counter() + self.last_at = self.started_at + self.total_tokens = 0 + self.rates: list[float] = [] + + def record(self, delta: str) -> None: + token_count = max(1, estimate_tokens(delta)) + now = time.perf_counter() + elapsed = max(now - self.last_at, 0.001) + self.total_tokens += token_count + self.rates.append(token_count / elapsed) + self.last_at = now + + def summary(self) -> dict[str, Any]: + elapsed = max(self.last_at - self.started_at, 0.001) + avg = self.total_tokens / elapsed if self.total_tokens else 0.0 + return { + "total_tokens": self.total_tokens, + "min_tokens_per_second": round(min(self.rates), 2) if self.rates else 0.0, + "avg_tokens_per_second": round(avg, 2), + "max_tokens_per_second": round(max(self.rates), 2) if self.rates else 0.0, + } + + class ChatRequest(BaseModel): message: str conversation_id: str | None = None @@ -297,6 +324,7 @@ def create_app() -> FastAPI: reasoning_parts: list[str] = [] content_parts: list[str] = [] + generation_stats = GenerationStats() try: messages = await runtime.context_builder.build_async_messages( task, history, memory_records @@ -367,12 +395,14 @@ def create_app() -> FastAPI: async for chunk in model_client.stream_chat("thinker", messages): delta = str(chunk.get("delta") or "") if chunk.get("type") == "reasoning_delta": + generation_stats.record(delta) reasoning_parts.append(delta) yield sse( "reasoning_delta", {"task_id": task.task_id, "delta": delta}, ) elif chunk.get("type") == "content_delta": + generation_stats.record(delta) content_parts.append(delta) yield sse( "content_delta", @@ -396,6 +426,7 @@ def create_app() -> FastAPI: { "role": "thinker", "model": model_client.get_role_config("thinker").model, + "generation_stats": generation_stats.summary(), }, ) await task_store.complete_task(task.task_id, content) @@ -420,6 +451,7 @@ def create_app() -> FastAPI: { "final_response": content, "reasoning_content": reasoning_content, + "generation_stats": generation_stats.summary(), }, ) yield sse( @@ -430,6 +462,7 @@ def create_app() -> FastAPI: "status": "completed", "final_response": content, "reasoning_content": reasoning_content, + "generation_stats": generation_stats.summary(), }, ) except asyncio.CancelledError: @@ -518,6 +551,7 @@ def create_app() -> FastAPI: async def generator(): reasoning_parts: list[str] = [] content_parts: list[str] = [] + generation_stats = GenerationStats() try: await task_store.update_status(task_id, "running") continued_event = await event_store.append( @@ -601,9 +635,11 @@ def create_app() -> FastAPI: async for chunk in model_client.stream_chat("thinker", messages): delta = str(chunk.get("delta") or "") if chunk.get("type") == "reasoning_delta": + generation_stats.record(delta) reasoning_parts.append(delta) yield sse("reasoning_delta", {"task_id": task_id, "delta": delta}) elif chunk.get("type") == "content_delta": + generation_stats.record(delta) content_parts.append(delta) yield sse("content_delta", {"task_id": task_id, "delta": delta}) @@ -624,6 +660,7 @@ def create_app() -> FastAPI: { "role": "thinker", "model": model_client.get_role_config("thinker").model, + "generation_stats": generation_stats.summary(), }, ) await task_store.complete_task(task_id, content) @@ -642,6 +679,7 @@ def create_app() -> FastAPI: { "final_response": content, "reasoning_content": reasoning_content, + "generation_stats": generation_stats.summary(), }, ) yield sse( @@ -652,6 +690,7 @@ def create_app() -> FastAPI: "status": "completed", "final_response": content, "reasoning_content": reasoning_content, + "generation_stats": generation_stats.summary(), }, ) except asyncio.CancelledError: @@ -702,6 +741,7 @@ def create_app() -> FastAPI: async def generator(): reasoning_parts: list[str] = [] content_parts: list[str] = [] + generation_stats = GenerationStats() try: await task_store.update_status(task_id, "running") continued_event = await event_store.append( @@ -762,9 +802,11 @@ def create_app() -> FastAPI: async for chunk in model_client.stream_chat("thinker", messages): delta = str(chunk.get("delta") or "") if chunk.get("type") == "reasoning_delta": + generation_stats.record(delta) reasoning_parts.append(delta) yield sse("reasoning_delta", {"task_id": task_id, "delta": delta}) elif chunk.get("type") == "content_delta": + generation_stats.record(delta) content_parts.append(delta) yield sse("content_delta", {"task_id": task_id, "delta": delta}) @@ -785,6 +827,7 @@ def create_app() -> FastAPI: { "role": "thinker", "model": model_client.get_role_config("thinker").model, + "generation_stats": generation_stats.summary(), }, ) await task_store.complete_task(task_id, content) @@ -803,6 +846,7 @@ def create_app() -> FastAPI: { "final_response": content, "reasoning_content": reasoning_content, + "generation_stats": generation_stats.summary(), }, ) yield sse( @@ -813,6 +857,7 @@ def create_app() -> FastAPI: "status": "completed", "final_response": content, "reasoning_content": reasoning_content, + "generation_stats": generation_stats.summary(), }, ) except asyncio.CancelledError: diff --git a/duck_core/tools/file_read.py b/duck_core/tools/file_read.py index d2a879f..6ba46bd 100644 --- a/duck_core/tools/file_read.py +++ b/duck_core/tools/file_read.py @@ -2,7 +2,7 @@ from pathlib import Path from typing import Any from duck_core.tools.base import ToolResult -from duck_core.tools.paths import WorkspacePathError, resolve_workspace_path +from duck_core.tools.paths import WorkspacePathError, candidate_path, resolve_workspace_path class FileReadTool: @@ -15,12 +15,14 @@ class FileReadTool: async def run(self, args: dict[str, Any]) -> ToolResult: raw_path = str(args.get("path", "")) + approved = bool(args.get("_approved")) try: - path = resolve_workspace_path(self.workspace, raw_path) + path = resolve_workspace_path(self.workspace, raw_path, allow_outside=approved) except WorkspacePathError as exc: - return ToolResult(ok=False, error=str(exc)) - if self._requires_approval(path): - return ToolResult(ok=False, error=f"Reading {raw_path} requires explicit approval") + path = candidate_path(self.workspace, raw_path) + return self._approval_required(raw_path, path, str(exc)) + if self._requires_approval(path) and not approved: + return self._approval_required(raw_path, path, f"Reading {raw_path} requires explicit approval") if not path.is_file(): return ToolResult(ok=False, error=f"File not found: {raw_path}") if path.stat().st_size > self.max_bytes: @@ -34,3 +36,15 @@ class FileReadTool: def _requires_approval(self, path: Path) -> bool: parts = set(path.parts) return path.name == ".env" or ".ssh" in parts or str(path) == "/etc/shadow" + + def _approval_required(self, raw_path: str, path: Path, reason: str) -> ToolResult: + return ToolResult( + ok=False, + error=reason, + metadata={ + "path": str(path), + "requires_approval": True, + "risk_level": self.risk_level, + "reason": reason, + }, + ) diff --git a/duck_core/tools/file_write.py b/duck_core/tools/file_write.py index 2dc1618..4c9799f 100644 --- a/duck_core/tools/file_write.py +++ b/duck_core/tools/file_write.py @@ -1,7 +1,7 @@ from typing import Any from duck_core.tools.base import ToolResult -from duck_core.tools.paths import WorkspacePathError, resolve_workspace_path +from duck_core.tools.paths import WorkspacePathError, candidate_path, resolve_workspace_path class FileWriteTool: @@ -15,15 +15,31 @@ class FileWriteTool: raw_path = str(args.get("path", "")) content = str(args.get("content", "")) overwrite = bool(args.get("overwrite", False)) + approved = bool(args.get("_approved")) try: - path = resolve_workspace_path(self.workspace, raw_path) + path = resolve_workspace_path(self.workspace, raw_path, allow_outside=approved) except WorkspacePathError as exc: - return ToolResult(ok=False, error=str(exc)) - if path.exists() and not overwrite: + path = candidate_path(self.workspace, raw_path) + return ToolResult( + ok=False, + error=f"{exc}. Writing outside workspace requires approval.", + metadata={ + "path": str(path), + "requires_approval": True, + "risk_level": self.risk_level, + "reason": "Path is outside workspace", + }, + ) + if path.exists() and not overwrite and not approved: return ToolResult( ok=False, error="Refusing to overwrite existing file without overwrite=true or approval", - metadata={"path": str(path)}, + metadata={ + "path": str(path), + "requires_approval": True, + "risk_level": self.risk_level, + "reason": "File exists and overwrite was not requested", + }, ) path.parent.mkdir(parents=True, exist_ok=True) existed = path.exists() diff --git a/duck_core/tools/list_dir.py b/duck_core/tools/list_dir.py index db9b504..efbefcd 100644 --- a/duck_core/tools/list_dir.py +++ b/duck_core/tools/list_dir.py @@ -1,7 +1,7 @@ from typing import Any from duck_core.tools.base import ToolResult -from duck_core.tools.paths import WorkspacePathError, resolve_workspace_path +from duck_core.tools.paths import WorkspacePathError, candidate_path, resolve_workspace_path class ListDirTool: @@ -14,10 +14,21 @@ class ListDirTool: async def run(self, args: dict[str, Any]) -> ToolResult: raw_path = str(args.get("path") or ".") + approved = bool(args.get("_approved")) try: - path = resolve_workspace_path(self.workspace, raw_path) + path = resolve_workspace_path(self.workspace, raw_path, allow_outside=approved) except WorkspacePathError as exc: - return ToolResult(ok=False, error=str(exc)) + path = candidate_path(self.workspace, raw_path) + return ToolResult( + ok=False, + error=f"{exc}. Access outside workspace requires approval.", + metadata={ + "path": str(path), + "requires_approval": True, + "risk_level": self.risk_level, + "reason": "Path is outside workspace", + }, + ) if not path.exists(): return ToolResult(ok=False, error=f"Directory not found: {raw_path}") if not path.is_dir(): diff --git a/duck_core/tools/paths.py b/duck_core/tools/paths.py index 9b8019d..53ccc72 100644 --- a/duck_core/tools/paths.py +++ b/duck_core/tools/paths.py @@ -5,9 +5,25 @@ class WorkspacePathError(ValueError): pass -def resolve_workspace_path(workspace: str, relative_path: str) -> Path: +def candidate_path(workspace: str, requested_path: str) -> Path: + root = Path(workspace).resolve() + path = Path(requested_path).expanduser() + if path.is_absolute(): + return path.resolve() + return (root / path).resolve() + + +def is_inside_workspace(workspace: str, path: Path) -> bool: root = Path(workspace).resolve() - path = (root / relative_path).resolve() if root != path and root not in path.parents: + return False + return True + + +def resolve_workspace_path( + workspace: str, relative_path: str, allow_outside: bool = False +) -> Path: + path = candidate_path(workspace, relative_path) + if not allow_outside and not is_inside_workspace(workspace, path): raise WorkspacePathError(f"Path escapes workspace: {relative_path}") return path diff --git a/duck_core/tools/search_files.py b/duck_core/tools/search_files.py index f3daaed..a051d82 100644 --- a/duck_core/tools/search_files.py +++ b/duck_core/tools/search_files.py @@ -2,7 +2,12 @@ from fnmatch import fnmatch from typing import Any from duck_core.tools.base import ToolResult -from duck_core.tools.paths import WorkspacePathError, resolve_workspace_path +from duck_core.tools.paths import ( + WorkspacePathError, + candidate_path, + is_inside_workspace, + resolve_workspace_path, +) class SearchFilesTool: @@ -22,11 +27,22 @@ class SearchFilesTool: max_matches = min(int(args.get("max_matches") or self.max_matches), self.max_matches) if not query: return ToolResult(ok=False, error="Search query is required") + approved = bool(args.get("_approved")) try: root = resolve_workspace_path(self.workspace, ".") - path = resolve_workspace_path(self.workspace, raw_path) + path = resolve_workspace_path(self.workspace, raw_path, allow_outside=approved) except WorkspacePathError as exc: - return ToolResult(ok=False, error=str(exc)) + path = candidate_path(self.workspace, raw_path) + return ToolResult( + ok=False, + error=f"{exc}. Searching outside workspace requires approval.", + metadata={ + "path": str(path), + "requires_approval": True, + "risk_level": self.risk_level, + "reason": "Path is outside workspace", + }, + ) if not path.exists(): return ToolResult(ok=False, error=f"Search path not found: {raw_path}") @@ -39,7 +55,10 @@ class SearchFilesTool: break if not candidate.is_file() or ".git" in candidate.parts: continue - relative = candidate.relative_to(root).as_posix() + if is_inside_workspace(self.workspace, candidate): + relative = candidate.relative_to(root).as_posix() + else: + relative = str(candidate) if not fnmatch(candidate.name, pattern) and not fnmatch(relative, pattern): continue if candidate.stat().st_size > self.max_file_bytes: diff --git a/duck_core/web/static/app.js b/duck_core/web/static/app.js index 438069a..d6556a5 100644 --- a/duck_core/web/static/app.js +++ b/duck_core/web/static/app.js @@ -369,6 +369,28 @@ function appendMessageText(article, delta) { document.querySelector("#messages").scrollTop = document.querySelector("#messages").scrollHeight; } +function formatTokenSpeed(value) { + const number = Number(value || 0); + return number.toFixed(number >= 10 ? 1 : 2); +} + +function setGenerationStats(article, stats) { + if (!article || !stats) return; + const bubble = article.querySelector(".bubble"); + if (!bubble) return; + let node = bubble.querySelector(".generation-stats"); + if (!node) { + node = document.createElement("div"); + node.className = "generation-stats"; + bubble.append(node); + } + node.textContent = [ + `min ${formatTokenSpeed(stats.min_tokens_per_second)} tok/s`, + `avg ${formatTokenSpeed(stats.avg_tokens_per_second)} tok/s`, + `max ${formatTokenSpeed(stats.max_tokens_per_second)} tok/s`, + ].join(" ยท "); +} + function appendInlineReasoning(article, delta) { const block = article?.querySelector(".message-reasoning"); const body = block?.querySelector("pre"); @@ -609,6 +631,7 @@ async function handleAssistantStreamEvent(pending, name, data, context) { setMessagePending(pending, data.final_response || "No final content returned."); } pending.querySelector(".message-meta span").textContent = data.status; + setGenerationStats(pending, data.generation_stats); setStatus("#task-status", data.task_id, data.status === "completed" ? "ok" : "warn"); finishInlineReasoning(pending, data.reasoning_content); await refreshEvents(data.task_id); diff --git a/duck_core/web/static/style.css b/duck_core/web/static/style.css index 6eef87b..9fea56a 100644 --- a/duck_core/web/static/style.css +++ b/duck_core/web/static/style.css @@ -700,6 +700,13 @@ dd { font-size: 13px; } +.generation-stats { + margin-top: 8px; + color: var(--muted); + font-size: 11px; + line-height: 1.4; +} + .debug-heading { display: flex; align-items: center; diff --git a/tests/smoke/test_api_stream_chat.py b/tests/smoke/test_api_stream_chat.py index 285a3f6..2c3d9a5 100644 --- a/tests/smoke/test_api_stream_chat.py +++ b/tests/smoke/test_api_stream_chat.py @@ -47,6 +47,10 @@ def test_stream_chat_endpoint_emits_sse_reasoning_and_content(tmp_path, monkeypa assert "event: reasoning_delta" in body assert "event: content_delta" in body assert "event: done" in body + assert '"generation_stats":' in body + assert '"min_tokens_per_second":' in body + assert '"avg_tokens_per_second":' in body + assert '"max_tokens_per_second":' in body assert "thinking" in body assert "answer" in body @@ -111,6 +115,55 @@ def test_stream_chat_endpoint_executes_tool_before_streaming_answer(tmp_path, mo assert "event: done" in body +def test_stream_chat_requests_approval_for_directory_outside_workspace(tmp_path, monkeypatch): + monkeypatch.setenv("DUCK_DB_PATH", str(tmp_path / "duck.sqlite3")) + workspace = tmp_path / "workspace" + outside = tmp_path / "outside" + workspace.mkdir() + outside.mkdir() + (outside / "note.txt").write_text("external") + + async def fake_chat(self, role, messages, temperature=None, max_output_tokens=None, response_format=None): + assert role == "action" + return ModelResponse( + role=role, + model="local-main", + content=json.dumps( + { + "kind": "action_directive", + "intent": "list external directory", + "risk_level": "low", + "actions": [ + { + "tool": "list_dir", + "args": {"path": str(outside)}, + "reason": "User asked for that directory", + } + ], + } + ), + reasoning_content=None, + raw={}, + latency_ms=1.0, + ) + + monkeypatch.setattr("duck_core.model_client.ModelClient.chat", fake_chat) + client = TestClient(create_app()) + + with client.stream( + "POST", + "/v1/chat/stream", + json={"message": "show outside", "workspace": str(workspace), "debug": True}, + ) as response: + body = "".join(response.iter_text()) + + assert response.status_code == 200 + assert "event: tool_approval_requested" in body + assert '"tool": "list_dir"' in body + assert str(outside) in body + assert '"status": "waiting_for_approval"' in body + + def test_continue_stream_executes_approved_tool_and_streams_answer(tmp_path, monkeypatch): monkeypatch.setenv("DUCK_DB_PATH", str(tmp_path / "duck.sqlite3")) action_calls = 0 diff --git a/tests/smoke/test_tool_gateway.py b/tests/smoke/test_tool_gateway.py index 13a1331..5a58dc5 100644 --- a/tests/smoke/test_tool_gateway.py +++ b/tests/smoke/test_tool_gateway.py @@ -102,6 +102,26 @@ async def test_tool_gateway_lists_workspace_directory(tmp_path): assert escaped.ok is False +@pytest.mark.asyncio +async def test_tool_gateway_requests_approval_for_directory_outside_workspace(tmp_path): + workspace = tmp_path / "workspace" + outside = tmp_path / "outside" + workspace.mkdir() + outside.mkdir() + (outside / "note.txt").write_text("external") + gateway = ToolGateway.default(str(workspace)) + action = {"tool": "list_dir", "args": {"path": str(outside)}} + + blocked = await gateway.run_action(action) + approved = await gateway.run_action(action, approved=True) + + assert blocked.ok is False + assert blocked.metadata["requires_approval"] is True + assert blocked.metadata["path"] == str(outside) + assert approved.ok is True + assert "note.txt" in approved.output + + @pytest.mark.asyncio async def test_tool_gateway_searches_file_contents(tmp_path): (tmp_path / "src").mkdir()