Skip to content
Open
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
60 changes: 44 additions & 16 deletions services/api/api/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,9 +100,23 @@

_ENGINE_HARNESSES = {"amp", "claude-code", "codex", "pi-mono"}
_REUSABLE_DB_STATES = {"running", "idle", "delivering", "error", "suspended"}
_CODEX_MODEL_PROFILES = frozenset({"fast", "think"})

IDLE_TTL_S = int(os.getenv("IDLE_TTL_S", "86400")) # 24 hours
SUSPENDED_RETENTION_S = int(os.getenv("SUSPENDED_RETENTION_S", str(7 * 24 * 60 * 60)))


def _runtime_model_for_engine(engine: str | None, model: str | None) -> str | None:
normalized = (model or "").strip().lower() or None
if not normalized:
return None
if engine == "codex":
return normalized
if normalized in _CODEX_MODEL_PROFILES:
return None
return normalized


MAX_ACTIVE_SANDBOX_SESSIONS = int(os.getenv("MAX_ACTIVE_SANDBOX_SESSIONS", "45"))
STREAM_EOF_REATTACH_MAX = int(os.getenv("STREAM_EOF_REATTACH_MAX", "6"))
STREAM_EOF_REATTACH_BACKOFF_S = float(os.getenv("STREAM_EOF_REATTACH_BACKOFF_S", "1.0"))
Expand Down Expand Up @@ -200,6 +214,7 @@ async def _db_get_session(thread_key: str) -> SandboxSession | None:
pool = _get_pool()
row = await pool.fetchrow(
"SELECT thread_key, sandbox_id, harness, engine, state, started_at, "
"model, "
"agent_thread_id, last_delivered_id, inflight_turn_id, inflight_turn_input, "
"inflight_attempts, last_result, trace_id "
"FROM sandbox_sessions WHERE thread_key = $1",
Expand All @@ -212,6 +227,7 @@ async def _db_get_session(thread_key: str) -> SandboxSession | None:
thread_key=row["thread_key"],
harness=row["harness"],
engine=row["engine"],
model=row["model"] or "",
started_at=row["started_at"].timestamp() if row["started_at"] else 0.0,
backend_name="kubernetes",
db_state=row["state"],
Expand Down Expand Up @@ -254,18 +270,19 @@ async def _db_insert_session(
session.trace_id = trace_id
row = await pool.fetchrow(
"INSERT INTO sandbox_sessions ("
"thread_key, sandbox_id, harness, engine, state, started_at, "
"thread_key, sandbox_id, harness, engine, model, state, started_at, "
"agent_thread_id, last_delivered_id, inflight_turn_id, inflight_turn_input, "
"inflight_started_at, inflight_attempts, last_result, last_result_at, trace_id"
") VALUES ($1, $2, $3, $4, $5, NOW(), $6, $7, $8::text, $9::jsonb, "
"CASE WHEN $8::text IS NULL THEN NULL ELSE NOW() END, $10, $11, "
"CASE WHEN $11::text = '' THEN NULL ELSE NOW() END, $12::uuid) "
") VALUES ($1, $2, $3, $4, $5, $6, NOW(), $7, $8, $9::text, $10::jsonb, "
"CASE WHEN $9::text IS NULL THEN NULL ELSE NOW() END, $11, $12, "
"CASE WHEN $12::text = '' THEN NULL ELSE NOW() END, $13::uuid) "
"ON CONFLICT (thread_key) DO NOTHING "
"RETURNING thread_key",
session.thread_key,
session.sandbox_id,
harness,
engine,
session.model or None,
initial_state,
agent_thread_id or None,
last_delivered_id or None,
Expand Down Expand Up @@ -836,7 +853,28 @@ async def get_or_spawn(
old_last_result: str = ""
old_trace_id: str = ""
pool = _get_pool()
effective_harness = harness or default_harness()
resolved_engine, resolved_persona, repo = _resolve_harness_profile(
effective_harness, persona=persona, engine_override=engine
)
requested_model = _runtime_model_for_engine(resolved_engine, model)

session = await _db_get_session(thread_key)
if session and requested_model and (session.model or "") != requested_model:
old_sandbox_id = session.sandbox_id
backend = get_backend()
with contextlib.suppress(Exception):
await backend.stop_by_id(old_sandbox_id)
await _db_delete_session(thread_key)
_drop_runtime(old_sandbox_id)
session = None
log.info(
"sandbox_replaced_for_model_profile",
thread_key=thread_key,
sandbox=old_sandbox_id[:12],
model=requested_model,
)

if session:
if session.db_state in _REUSABLE_DB_STATES:
backend = get_backend()
Expand Down Expand Up @@ -913,17 +951,10 @@ async def get_or_spawn(

thread_trace_id = await get_or_create_thread_trace_id(pool, thread_key)

effective_harness = harness or default_harness()

# Resolve harness profile (engine, persona, repo) once for both warm and cold paths
resolved_engine, resolved_persona, repo = _resolve_harness_profile(
effective_harness, persona=persona, engine_override=engine
)

# Try warm pool first
should_try_warm = (
not engine
and not model
and not requested_model
and not old_agent_thread_id
and not old_inflight_turn_id
and not (effective_harness == "amp" and resolved_engine == "codex")
Expand Down Expand Up @@ -958,9 +989,6 @@ async def get_or_spawn(
return claimed

# Cold spawn
resolved_engine, resolved_persona, repo = _resolve_harness_profile(
effective_harness, persona=persona, engine_override=engine
)
backend = get_backend()
await _evict_idle_sessions_for_capacity(backend)
trace_id = old_trace_id or thread_trace_id or str(uuid.uuid4())
Expand All @@ -970,7 +998,7 @@ async def get_or_spawn(
resolved_engine,
persona=resolved_persona,
repo=repo,
model=model,
model=requested_model,
resume_thread_id=old_agent_thread_id or None,
trace_id=trace_id,
)
Expand Down
49 changes: 39 additions & 10 deletions services/api/api/runtime_control.py
Original file line number Diff line number Diff line change
Expand Up @@ -303,6 +303,7 @@ def _agent_session_title(
# ── Per-message header (rendered italic at the top of every assistant message) ──

_DEFAULT_CLAUDE_MODEL = "claude-opus-4-8"
_DEFAULT_CODEX_MODEL_PROFILE = "fast"

_CLAUDE_MODEL_ALIASES: dict[str, str] = {
"opus": _DEFAULT_CLAUDE_MODEL,
Expand All @@ -327,6 +328,19 @@ def _resolve_codex_model_label(model: str | None) -> str:
return f"codex-{raw}"


def _default_model_for_assignment(
*,
harness: str | None,
engine: str | None,
model: str | None,
) -> str | None:
if model:
return model
if (engine or harness or "").strip().lower() == "codex":
return _DEFAULT_CODEX_MODEL_PROFILE
return None


def _engine_model_label(
*,
engine: str | None,
Expand Down Expand Up @@ -576,7 +590,7 @@ async def _write_agents_override(runtime_id: str, agents_md_override: str) -> No

async def get_active_assignment(pool, thread_key: str) -> dict[str, Any] | None:
row = await pool.fetchrow(
"SELECT thread_key, assignment_generation, runtime_id, harness, engine, persona_id, "
"SELECT thread_key, assignment_generation, runtime_id, harness, engine, model, persona_id, "
"prompt_ref, effective_agents_md_sha256, agents_md_override, state "
"FROM agent_runtime_assignments "
"WHERE thread_key = $1 AND state = 'active' "
Expand Down Expand Up @@ -636,7 +650,11 @@ async def spawn_assignment(
effective_harness = active_assignment.get("harness") or default_harness()
effective_engine = active_assignment.get("engine")
effective_persona_id = active_assignment.get("persona_id")
effective_model = None
effective_model = _default_model_for_assignment(
harness=effective_harness,
engine=effective_engine,
model=model,
)
effective_agents_md_override = active_assignment.get("agents_md_override")
else:
# Explicit harness wins; otherwise inherit from the persona's declared
Expand All @@ -649,7 +667,11 @@ async def spawn_assignment(
effective_harness = default_harness()
effective_engine = engine
effective_persona_id = persona_id
effective_model = model
effective_model = _default_model_for_assignment(
harness=effective_harness,
engine=effective_engine,
model=model,
)
effective_agents_md_override = agents_md_override

payload = {
Expand Down Expand Up @@ -738,7 +760,7 @@ async def spawn_assignment(
return decode_jsonb(existing_idem["response_json"], {})

active = await conn.fetchrow(
"SELECT assignment_generation, runtime_id, harness, engine, persona_id, "
"SELECT assignment_generation, runtime_id, harness, engine, model, persona_id, "
"prompt_ref, effective_agents_md_sha256, agents_md_override "
"FROM agent_runtime_assignments "
"WHERE thread_key = $1 AND state = 'active' "
Expand All @@ -758,11 +780,15 @@ async def spawn_assignment(
)
generation = int(active["assignment_generation"])
runtime_id = active["runtime_id"]
if runtime_id != session.sandbox_id:
if (
runtime_id != session.sandbox_id
or (active["model"] or None) != (session.model or None)
):
await conn.execute(
"UPDATE agent_runtime_assignments SET runtime_id = $1, updated_at = NOW() "
"WHERE thread_key = $2 AND assignment_generation = $3",
"UPDATE agent_runtime_assignments SET runtime_id = $1, model = $2, updated_at = NOW() "
"WHERE thread_key = $3 AND assignment_generation = $4",
session.sandbox_id,
session.model or None,
thread_key,
generation,
)
Expand All @@ -784,14 +810,15 @@ async def spawn_assignment(
)
await conn.execute(
"INSERT INTO agent_runtime_assignments ("
"thread_key, assignment_generation, runtime_id, harness, engine, "
"thread_key, assignment_generation, runtime_id, harness, engine, model, "
"persona_id, prompt_ref, effective_agents_md_sha256, agents_md_override, state"
") VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, 'active')",
") VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, 'active')",
thread_key,
generation,
session.sandbox_id,
session.harness,
session.engine,
session.model or None,
effective_persona_id,
prompt_ref,
prompt_sha,
Expand All @@ -811,6 +838,7 @@ async def spawn_assignment(
"assignment_state": assignment_state,
"assignment_generation": generation,
"persona_id": resolved_persona,
"model": session.model or None,
"prompt_ref": resolved_prompt_ref,
"effective_agents_md_sha256": resolved_prompt_sha,
}
Expand Down Expand Up @@ -2797,7 +2825,7 @@ async def _process_execution_impl(pool, row: dict[str, Any]) -> None:
)

assignment = await pool.fetchrow(
"SELECT harness, engine, runtime_id, agents_md_override, persona_id, prompt_ref, effective_agents_md_sha256 "
"SELECT harness, engine, model, runtime_id, agents_md_override, persona_id, prompt_ref, effective_agents_md_sha256 "
"FROM agent_runtime_assignments "
"WHERE thread_key = $1 AND assignment_generation = $2",
thread_key,
Expand Down Expand Up @@ -2890,6 +2918,7 @@ async def _process_execution_impl(pool, row: dict[str, Any]) -> None:
assignment["harness"],
engine=assignment["engine"],
persona=assignment["persona_id"],
model=assignment["model"],
)
if session.sandbox_id != assignment["runtime_id"]:
await pool.execute(
Expand Down
1 change: 1 addition & 0 deletions services/api/api/sandbox/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ class SandboxSession:
thread_key: str
harness: str
engine: str
model: str = ""
started_at: float = 0.0
backend_name: str = "" # "kubernetes", etc.
db_state: str = ""
Expand Down
6 changes: 4 additions & 2 deletions services/api/api/sandbox/kubernetes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1498,9 +1498,10 @@ async def create(
if engine == "claude-code" and model:
env.append(f"CLAUDE_MODEL={model}")
if engine == "codex" and model:
if model != "fast":
if model not in {"fast", "think"}:
raise ValueError(f"unknown Codex model profile: {model}")
env.append("CODEX_MODEL_PROFILE=fast")
if model == "fast":
env.append("CODEX_MODEL_PROFILE=fast")
if engine == "claude-code" and resume_thread_id:
env.append(f"CLAUDE_CONTINUE_SESSION_ID={resume_thread_id}")
if persona:
Expand Down Expand Up @@ -1722,6 +1723,7 @@ async def create(
thread_key=thread_key,
harness=harness,
engine=engine,
model=model or "",
started_at=time.time(),
backend_name=self.name,
trace_id=trace_id or "",
Expand Down
5 changes: 4 additions & 1 deletion services/api/api/workflow_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -993,19 +993,22 @@ async def _compute_agent_session_header(

persona = selector.get("persona_id")
harness = selector.get("harness")
model = selector.get("model")
engine: str | None = None
if not persona or not harness:
active = await get_active_assignment(pool, thread_key)
if isinstance(active, dict):
persona = persona or _nonempty(active.get("persona_id"))
harness = harness or _nonempty(active.get("harness"))
engine = _nonempty(active.get("engine"))
model = model or _nonempty(active.get("model"))
if persona and not engine:
engine = _persona_default_engine(persona)
return _agent_session_header(
persona_id=persona,
engine=engine,
harness=harness,
model=model,
)


Expand Down Expand Up @@ -1216,7 +1219,7 @@ async def _dispatch() -> dict[str, Any]:
else:
effective_delivery = dict(run_in.get("delivery") or {})
effective_history = history_messages or run_in.get("history_messages") or []
selector = {"persona_id": persona, "harness": harness}
selector = {"persona_id": persona, "harness": harness, "model": model}
slackbot_session_id: str | None = None

try:
Expand Down
10 changes: 7 additions & 3 deletions services/api/api/workflows/slack_thread_turn.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,10 @@
"claude": "claude-code",
"pi": "pi-mono",
}
_MODEL_PROFILE_FLAGS = frozenset({"fast"})
_MODEL_PROFILE_FLAGS = {
"fast": "fast",
"think": "think",
}
_PROMPT_FLAG_SKIP = frozenset({"engine", "model", "opus", "sonnet", "haiku"})
_PROMPT_FLAG_VALUE_SKIP = frozenset({"engine", "model"})
_PROMPT_FLAG_RE = re.compile(
Expand Down Expand Up @@ -68,7 +71,8 @@ class PromptSelection:

Both fields are optional and orthogonal: ``--invest`` sets only
``persona``, ``--claude`` sets only ``harness``, and ``--invest --claude``
sets both. ``--fast`` selects the fast Codex model profile for the sandbox
sets both. Slack turns default to the fast Codex model profile; ``--think``
opts into the deployment's higher-reasoning Codex profile for the sandbox
spawned for this turn. The downstream resolver applies ``harness`` as the
engine override and ``persona`` as the system-prompt overlay.
"""
Expand Down Expand Up @@ -140,7 +144,7 @@ def _classify_flag(
if resolved in personas or flag in personas:
return None, resolved, None
if resolved in _MODEL_PROFILE_FLAGS:
return None, None, resolved
return None, None, _MODEL_PROFILE_FLAGS[resolved]
return None, None, None


Expand Down
15 changes: 15 additions & 0 deletions services/api/db/migrations/041_track_runtime_model_profiles.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
-- migrate:up

ALTER TABLE sandbox_sessions
ADD COLUMN IF NOT EXISTS model TEXT;

ALTER TABLE agent_runtime_assignments
ADD COLUMN IF NOT EXISTS model TEXT;

-- migrate:down

ALTER TABLE agent_runtime_assignments
DROP COLUMN IF EXISTS model;

ALTER TABLE sandbox_sessions
DROP COLUMN IF EXISTS model;
Loading
Loading