Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 119 additions & 1 deletion tests/test_routes_project_canvas.py
Original file line number Diff line number Diff line change
Expand Up @@ -573,7 +573,6 @@ async def test_read_without_flag_is_403(self, ctx):
)
assert resp.status_code == 403


@pytest.mark.asyncio
class TestAgentCanvasWriteGating:
async def test_write_allowed_with_scope_and_flag(self, ctx):
Expand Down Expand Up @@ -790,3 +789,122 @@ async def test_patch_rejects_agent_token(self, ctx):
)
assert resp.status_code == 401


# ---------------------------------------------------------------------------
# Slice 5: payload size cap (413) + agent write rate limit (429).
# ---------------------------------------------------------------------------


class TestCanvasPayloadSizeCap:
"""The 64 KB payload cap applies to all principals (agents and humans)."""

@pytest.mark.asyncio
async def test_create_oversized_payload_returns_413(self, ctx):
pid = await _new_project(ctx, "payload-cap")
oversized = "x" * (65 * 1024)
body = {
"kind": "text",
"x": 0, "y": 0, "w": 100, "h": 50,
"payload": {"text": oversized},
}
resp = await ctx.client.post(
f"/api/projects/{pid}/canvas/elements", json=body,
)
assert resp.status_code == 413, resp.text

@pytest.mark.asyncio
async def test_patch_oversized_payload_returns_413(self, ctx):
pid = await _new_project(ctx, "payload-cap-patch")
el = await _create_note(ctx.client, pid)
oversized = "x" * (65 * 1024)
resp = await ctx.client.patch(
f"/api/projects/{pid}/canvas/elements/{el['id']}",
json={"payload": {"text": oversized}},
)
assert resp.status_code == 413, resp.text

@pytest.mark.asyncio
async def test_create_normal_payload_succeeds(self, ctx):
pid = await _new_project(ctx, "payload-cap-ok")
body = {
"kind": "text",
"x": 0, "y": 0, "w": 100, "h": 50,
"payload": {"text": "normal content"},
}
resp = await ctx.client.post(
f"/api/projects/{pid}/canvas/elements", json=body,
)
assert resp.status_code == 201, resp.text


class TestAgentWriteRateLimit:
"""Agents are throttled to 30 canvas writes per 60 s rolling window.
Humans (session principals) are never throttled."""

@pytest.mark.asyncio
async def test_agent_exceeding_window_returns_429(self, ctx):
import time
import tinyagentos.routes.project_canvas as canvas_mod

pid = await _new_project(ctx, "rate-limit-agent")
cid, token = await _mint_agent(ctx, pid, ("canvas_write",))
await _add_member(ctx, pid, cid)
await _grant_canvas(ctx, pid, cid, edit=True)

limiter = canvas_mod._canvas_write_limiter
max_attempts = canvas_mod._CANVAS_WRITE_MAX_ATTEMPTS
now = time.monotonic()
with limiter._lock:
limiter._log[cid] = [now - 1.0] * max_attempts

async with _bare(ctx.app) as bare:
resp = await bare.post(
f"/api/projects/{pid}/canvas/elements",
json={"kind": "note", "x": 0, "y": 0, "w": 1, "h": 1,
"payload": {"text": "rate-limited"}},
headers=_hdr(token),
)
assert resp.status_code == 429, resp.text

@pytest.mark.asyncio
async def test_agent_delete_exceeding_window_returns_429(self, ctx):
import time
import tinyagentos.routes.project_canvas as canvas_mod

pid = await _new_project(ctx, "rate-limit-agent-del")
cid, token = await _mint_agent(ctx, pid, ("canvas_write",))
await _add_member(ctx, pid, cid)
await _grant_canvas(ctx, pid, cid, edit=True)

created = await ctx.client.post(
f"/api/projects/{pid}/canvas/elements",
json={"kind": "note", "x": 0, "y": 0, "w": 1, "h": 1,
"payload": {"text": "to-delete"}},
)
assert created.status_code == 201
eid = created.json()["element"]["id"]

limiter = canvas_mod._canvas_write_limiter
max_attempts = canvas_mod._CANVAS_WRITE_MAX_ATTEMPTS
now = time.monotonic()
with limiter._lock:
limiter._log[cid] = [now - 1.0] * max_attempts

async with _bare(ctx.app) as bare:
resp = await bare.delete(
f"/api/projects/{pid}/canvas/elements/{eid}",
headers=_hdr(token),
)
assert resp.status_code == 429, resp.text

@pytest.mark.asyncio
async def test_human_write_not_rate_limited(self, ctx):
pid = await _new_project(ctx, "rate-limit-human")
for _ in range(31):
resp = await ctx.client.post(
f"/api/projects/{pid}/canvas/elements",
json={"kind": "note", "x": 0, "y": 0, "w": 1, "h": 1,
"payload": {"text": "human-write"}},
)
assert resp.status_code == 201, resp.text

86 changes: 82 additions & 4 deletions tinyagentos/routes/project_canvas.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@
import asyncio
import json
import logging
import threading
import time
from collections import OrderedDict
from typing import Literal

from fastapi import APIRouter, HTTPException, Request
Expand All @@ -20,12 +23,63 @@
logger = logging.getLogger(__name__)
router = APIRouter()

# Agent scopes (slice 1) that unlock canvas read / write on a project the token
# is bound to. The write scope is strictly narrower than read: holding
# canvas_write must NEVER satisfy a read, and vice versa (D3 enforcement matrix).
_CANVAS_READ_SCOPE = "canvas_read"
_CANVAS_WRITE_SCOPE = "canvas_write"

_CANVAS_PAYLOAD_MAX_BYTES = 64 * 1024
_CANVAS_WRITE_MAX_ATTEMPTS = 30
_CANVAS_WRITE_WINDOW_SECONDS = 60
_CANVAS_WRITE_MAX_KEYS = 10_000


class _CanvasWriteLimiter:
def __init__(self, max_attempts: int = _CANVAS_WRITE_MAX_ATTEMPTS,
window_seconds: int = _CANVAS_WRITE_WINDOW_SECONDS):
self._max = max_attempts
self._window = window_seconds
self._log: OrderedDict[str, list[float]] = OrderedDict()
self._lock = threading.Lock()

def _prune(self, key: str) -> None:
cutoff = time.monotonic() - self._window
if key not in self._log:
return
self._log[key] = [t for t in self._log[key] if t > cutoff]
if not self._log[key]:
del self._log[key]
else:
self._log.move_to_end(key)

def _ensure_capacity(self) -> None:
while len(self._log) >= _CANVAS_WRITE_MAX_KEYS:
self._log.popitem(last=False)

def is_limited(self, key: str) -> bool:
with self._lock:
self._prune(key)
return len(self._log.get(key, [])) >= self._max

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WARNING: Non-atomic check-then-act allows the rate limit to be exceeded under concurrency

The limit decision here and the increment in record are separate locked operations, but record is only called after the awaited store write (await cs.add_element / update_element / delete_element). Under a burst, more than 30 concurrent requests can all observe count < 30 from is_limited before any of them calls record, so all of them succeed and the 30/60s budget is exceeded.

Fix: combine the check and increment into a single atomic method (acquire lock once, if count >= max return limited without recording, else append timestamp and return allowed) and call it exactly once per request.


Reply with @kilocode-bot fix it to have Kilo Code address this issue.


def record(self, key: str) -> None:
with self._lock:
self._prune(key)
if key not in self._log:
self._ensure_capacity()
self._log[key] = []
self._log[key].append(time.monotonic())
self._log.move_to_end(key)


_canvas_write_limiter = _CanvasWriteLimiter()

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SUGGESTION: In-memory singleton limiter: per-process, non-persistent, and global per agent

The limiter is a module-level in-memory dict, so (a) it resets on every process restart, (b) it is not shared across multiple workers/instances — so the effective limit becomes 30 * workers per window — and (c) it keys only on actor_id, meaning one agent's 30/60s budget is shared across all projects it can write to rather than being scoped per (agent, project).

If cross-project isolation or multi-worker correctness matters, back this with a shared store (Redis/DB) and key on (actor_id, project_id). Also consider adding a Retry-After header on the 429 responses for correct client backoff.


Reply with @kilocode-bot fix it to have Kilo Code address this issue.



def _check_payload_size(payload: dict) -> None:
size = len(json.dumps(payload, default=str).encode("utf-8"))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WARNING: Payload cap is enforced after full request parsing and may mis-measure

_check_payload_size runs on the already-deserialized pydantic payload, i.e. only after FastAPI/pydantic have fully materialized the entire JSON request body in memory. A 64 KB payload cap therefore does not protect against a multi-MB request body (memory/DoS): the oversized body is already parsed before this guard returns 413. Also json.dumps(payload, default=str) may over/under-estimate size versus the store's actual serialization if payload ever holds non-JSON-native objects.

For real protection, enforce a raw request-body size limit (ASGI/gateway layer or a streaming Request body read) before deserialization.


Reply with @kilocode-bot fix it to have Kilo Code address this issue.

if size > _CANVAS_PAYLOAD_MAX_BYTES:
raise HTTPException(
status_code=413,
detail=f"canvas payload exceeds {_CANVAS_PAYLOAD_MAX_BYTES // 1024} KB limit",
)


def _user_id(request: Request) -> str:
uid = getattr(request.state, "user_id", None)
Expand Down Expand Up @@ -136,8 +190,14 @@ async def create_canvas_element(
if isinstance(auth, JSONResponse):
return auth
actor_kind, actor_id = auth
cs = request.app.state.project_canvas_store
if actor_kind == "agent" and _canvas_write_limiter.is_limited(actor_id):
return JSONResponse(
{"error": "rate_limit", "message": "too many canvas writes, try again later"},
status_code=429,
)
element = payload.model_dump()
_check_payload_size(element.get("payload") or {})
cs = request.app.state.project_canvas_store
element_id = element.pop("element_id", None)
if element["kind"] == "link":
url = (element.get("payload") or {}).get("url")
Expand All @@ -153,6 +213,8 @@ async def create_canvas_element(
)
except ValueError as e:
return JSONResponse({"error": str(e)}, status_code=400)
if actor_kind == "agent":
_canvas_write_limiter.record(actor_id)
return {"element": new_el}


Expand All @@ -174,8 +236,15 @@ async def update_canvas_element(
if isinstance(auth, JSONResponse):
return auth
actor_kind, actor_id = auth
if actor_kind == "agent" and _canvas_write_limiter.is_limited(actor_id):
return JSONResponse(
{"error": "rate_limit", "message": "too many canvas writes, try again later"},
status_code=429,
)
cs = request.app.state.project_canvas_store
patch = {k: v for k, v in payload.model_dump().items() if v is not None}
if "payload" in patch:
_check_payload_size(patch["payload"])
try:
updated = await cs.update_element(
project_id=project_id, element_id=element_id, patch=patch,
Expand All @@ -185,6 +254,8 @@ async def update_canvas_element(
return JSONResponse({"error": "permission_denied", "message": str(e)}, status_code=403)
except ValueError as e:
return JSONResponse({"error": str(e)}, status_code=404)
if actor_kind == "agent":
_canvas_write_limiter.record(actor_id)
return {"element": updated}


Expand All @@ -194,6 +265,11 @@ async def delete_canvas_element(project_id: str, element_id: str, request: Reque
if isinstance(auth, JSONResponse):
return auth
actor_kind, actor_id = auth
if actor_kind == "agent" and _canvas_write_limiter.is_limited(actor_id):
return JSONResponse(
{"error": "rate_limit", "message": "too many canvas writes, try again later"},
status_code=429,
)
cs = request.app.state.project_canvas_store
try:
await cs.delete_element(
Expand All @@ -202,6 +278,8 @@ async def delete_canvas_element(project_id: str, element_id: str, request: Reque
)
except CanvasPermissionError as e:
return JSONResponse({"error": "permission_denied", "message": str(e)}, status_code=403)
if actor_kind == "agent":
_canvas_write_limiter.record(actor_id)
return Response(status_code=204)


Expand Down