wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Isolate plan and task runtime per session

Browse Source
This commit is contained in:
wuyang6
2026-05-13 20:07:09 +08:00
parent f2989cd48b
commit 0eb57f3f42
8 changed files with 146 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent_context.py
+8 -6
View File
@@ -258,12 +258,14 @@ def _get_user_context_cached(
lsp_runtime = LSPRuntime.from_workspace(Path(cwd), additional_working_directories)
if lsp_runtime.has_lsp_support():
context['lspRuntime'] = lsp_runtime.render_summary()
plan_runtime = PlanRuntime.from_workspace(Path(cwd))
if plan_runtime.steps:
context['planRuntime'] = plan_runtime.render_summary()
task_runtime = TaskRuntime.from_workspace(Path(cwd))
if task_runtime.tasks:
context['taskRuntime'] = task_runtime.render_summary()
if scratchpad_directory:
scratchpad_path = Path(scratchpad_directory)
plan_runtime = PlanRuntime.from_storage_path(scratchpad_path / 'plan_runtime.json')
if plan_runtime.steps:
context['planRuntime'] = plan_runtime.render_summary()
task_runtime = TaskRuntime.from_storage_path(scratchpad_path / 'task_runtime.json')
if task_runtime.tasks:
context['taskRuntime'] = task_runtime.render_summary()
team_runtime = TeamRuntime.from_workspace(Path(cwd), additional_working_directories)
if team_runtime.has_team_state():
context['teamRuntime'] = team_runtime.render_summary()
src/agent_runtime.py
+68 -6
View File
@@ -245,9 +245,9 @@ class LocalCodingAgent:
tuple(str(path) for path in self.runtime_config.additional_working_directories),
)
if self.plan_runtime is None:
self.plan_runtime = PlanRuntime.from_workspace(self.runtime_config.cwd)
self.plan_runtime = PlanRuntime()
if self.task_runtime is None:
self.task_runtime = TaskRuntime.from_workspace(self.runtime_config.cwd)
self.task_runtime = TaskRuntime()
if self.team_runtime is None:
self.team_runtime = TeamRuntime.from_workspace(
self.runtime_config.cwd,
@@ -424,6 +424,7 @@ class LocalCodingAgent:
self.plugin_runtime.restore_session_state({})
session_id = session_id or uuid4().hex
scratchpad_directory = self._ensure_scratchpad_directory(session_id)
self._bind_session_plan_task_runtime(scratchpad_directory)
result = self._run_prompt(
prompt,
base_session=None,
@@ -472,6 +473,7 @@ class LocalCodingAgent:
if stored_session.scratchpad_directory
else self._ensure_scratchpad_directory(stored_session.session_id)
)
self._bind_session_plan_task_runtime(scratchpad_directory)
result = self._run_prompt(
prompt,
base_session=session,
@@ -544,6 +546,8 @@ class LocalCodingAgent:
self.tool_context = replace(
self.tool_context,
scratchpad_directory=scratchpad_directory,
plan_runtime=self.plan_runtime,
task_runtime=self.task_runtime,
)
tool_specs = [tool.to_openai_tool() for tool in self.tool_registry.values()]
starting_usage = UsageStats()
@@ -3560,6 +3564,36 @@ class LocalCodingAgent:
scratchpad_directory.mkdir(parents=True, exist_ok=True)
return scratchpad_directory
def _bind_session_plan_task_runtime(self, scratchpad_directory: Path | None) -> None:
if scratchpad_directory is None:
self.plan_runtime = PlanRuntime()
self.task_runtime = TaskRuntime()
else:
self.plan_runtime = PlanRuntime.from_storage_path(
scratchpad_directory / 'plan_runtime.json'
)
self.task_runtime = TaskRuntime.from_storage_path(
scratchpad_directory / 'task_runtime.json'
)
self.tool_context = replace(
self.tool_context,
plan_runtime=self.plan_runtime,
task_runtime=self.task_runtime,
scratchpad_directory=scratchpad_directory,
)
def _session_plan_runtime_path(self) -> Path:
scratchpad_directory = self.tool_context.scratchpad_directory
if scratchpad_directory is None:
return PlanRuntime().storage_path
return scratchpad_directory / 'plan_runtime.json'
def _session_task_runtime_path(self) -> Path:
scratchpad_directory = self.tool_context.scratchpad_directory
if scratchpad_directory is None:
return TaskRuntime().storage_path
return scratchpad_directory / 'task_runtime.json'
def _append_file_history_replay_if_needed(
self,
session: AgentSessionState,
@@ -4833,9 +4867,25 @@ class LocalCodingAgent:
if tool_name == 'config_set':
self.config_runtime = ConfigRuntime.from_workspace(self.runtime_config.cwd)
if tool_name.startswith('task_') or tool_name == 'todo_write':
self.task_runtime = TaskRuntime.from_workspace(self.runtime_config.cwd)
self.task_runtime = TaskRuntime.from_storage_path(
(
self.task_runtime.storage_path
if self.task_runtime is not None
else self._session_task_runtime_path()
)
)
if tool_name.startswith('plan_') or tool_name == 'update_plan':
self.plan_runtime = PlanRuntime.from_workspace(self.runtime_config.cwd)
self.plan_runtime = PlanRuntime.from_storage_path(
(
self.plan_runtime.storage_path
if self.plan_runtime is not None
else self._session_plan_runtime_path()
)
)
if self.task_runtime is not None:
self.task_runtime = TaskRuntime.from_storage_path(
self.task_runtime.storage_path
)
if tool_name.startswith('team_') or tool_name == 'send_message':
self.team_runtime = TeamRuntime.from_workspace(
self.runtime_config.cwd,
@@ -4911,8 +4961,20 @@ class LocalCodingAgent:
self.runtime_config.cwd,
additional_dirs,
)
self.task_runtime = TaskRuntime.from_workspace(self.runtime_config.cwd)
self.plan_runtime = PlanRuntime.from_workspace(self.runtime_config.cwd)
self.task_runtime = TaskRuntime.from_storage_path(
(
self.task_runtime.storage_path
if self.task_runtime is not None
else self._session_task_runtime_path()
)
)
self.plan_runtime = PlanRuntime.from_storage_path(
(
self.plan_runtime.storage_path
if self.plan_runtime is not None
else self._session_plan_runtime_path()
)
)
self.team_runtime = TeamRuntime.from_workspace(
self.runtime_config.cwd,
additional_dirs,
src/plan_runtime.py
+5
View File
@@ -106,6 +106,11 @@ class PlanRuntime:
@classmethod
def from_workspace(cls, cwd: Path) -> 'PlanRuntime':
storage_path = (cwd.resolve() / DEFAULT_PLAN_RUNTIME_PATH).resolve()
return cls.from_storage_path(storage_path)
@classmethod
def from_storage_path(cls, storage_path: Path) -> 'PlanRuntime':
storage_path = storage_path.resolve()
if not storage_path.exists():
return cls(storage_path=storage_path)
try:
src/task_runtime.py
+5
View File
@@ -36,6 +36,11 @@ class TaskRuntime:
@classmethod
def from_workspace(cls, cwd: Path) -> 'TaskRuntime':
storage_path = (cwd.resolve() / DEFAULT_TASK_RUNTIME_PATH).resolve()
return cls.from_storage_path(storage_path)
@classmethod
def from_storage_path(cls, storage_path: Path) -> 'TaskRuntime':
storage_path = storage_path.resolve()
if not storage_path.exists():
return cls(tasks=(), storage_path=storage_path)
try:
tests/test_agent_context.py
+34 -4
View File
@@ -199,10 +199,15 @@ class AgentContextTests(unittest.TestCase):
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir) / 'repo'
workspace.mkdir(parents=True)
runtime = TaskRuntime.from_workspace(workspace)
scratchpad = workspace / '.port_sessions' / 'accounts' / 'alice' / 'sessions' / 's1' / 'scratchpad'
scratchpad.mkdir(parents=True)
runtime = TaskRuntime.from_storage_path(scratchpad / 'task_runtime.json')
runtime.create_task(title='Review task runtime')
snapshot = build_context_snapshot(AgentRuntimeConfig(cwd=workspace))
snapshot = build_context_snapshot(
AgentRuntimeConfig(cwd=workspace),
scratchpad_directory=scratchpad,
)
self.assertIn('taskRuntime', snapshot.user_context)
self.assertIn('Total tasks: 1', snapshot.user_context['taskRuntime'])
@@ -211,17 +216,42 @@ class AgentContextTests(unittest.TestCase):
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir) / 'repo'
workspace.mkdir(parents=True)
plan_runtime = PlanRuntime.from_workspace(workspace)
scratchpad = workspace / '.port_sessions' / 'accounts' / 'alice' / 'sessions' / 's1' / 'scratchpad'
scratchpad.mkdir(parents=True)
plan_runtime = PlanRuntime.from_storage_path(scratchpad / 'plan_runtime.json')
plan_runtime.update_plan(
[{'step': 'Inspect the runtime', 'status': 'in_progress'}],
explanation='Use a stored plan.',
)
snapshot = build_context_snapshot(AgentRuntimeConfig(cwd=workspace))
snapshot = build_context_snapshot(
AgentRuntimeConfig(cwd=workspace),
scratchpad_directory=scratchpad,
)
self.assertIn('planRuntime', snapshot.user_context)
self.assertIn('Total plan steps: 1', snapshot.user_context['planRuntime'])
def test_user_context_ignores_workspace_plan_task_when_session_scratchpad_is_set(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir) / 'repo'
workspace.mkdir(parents=True)
TaskRuntime.from_workspace(workspace).create_task(title='Global task')
PlanRuntime.from_workspace(workspace).update_plan(
[{'step': 'Global plan', 'status': 'in_progress'}],
explanation='This must not leak into a session.',
)
scratchpad = workspace / '.port_sessions' / 'accounts' / 'alice' / 'sessions' / 's1' / 'scratchpad'
scratchpad.mkdir(parents=True)
snapshot = build_context_snapshot(
AgentRuntimeConfig(cwd=workspace),
scratchpad_directory=scratchpad,
)
self.assertNotIn('taskRuntime', snapshot.user_context)
self.assertNotIn('planRuntime', snapshot.user_context)
def test_user_context_loads_team_runtime_summary(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir) / 'repo'
tests/test_agent_prompting.py
+16 -4
View File
@@ -269,11 +269,17 @@ class AgentPromptingTests(unittest.TestCase):
def test_prompt_builder_mentions_tasks_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
runtime = TaskRuntime.from_workspace(workspace)
scratchpad = workspace / '.port_sessions' / 'accounts' / 'alice' / 'sessions' / 's1' / 'scratchpad'
scratchpad.mkdir(parents=True)
runtime = TaskRuntime.from_storage_path(scratchpad / 'task_runtime.json')
runtime.create_task(title='Inspect runtime tasks')
runtime_config = AgentRuntimeConfig(cwd=workspace)
model_config = ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct')
prompt_context = build_prompt_context(runtime_config, model_config)
prompt_context = build_prompt_context(
runtime_config,
model_config,
scratchpad_directory=scratchpad,
)
parts = build_system_prompt_parts(
prompt_context=prompt_context,
runtime_config=runtime_config,
@@ -305,14 +311,20 @@ class AgentPromptingTests(unittest.TestCase):
def test_prompt_builder_mentions_planning_when_runtime_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
runtime = PlanRuntime.from_workspace(workspace)
scratchpad = workspace / '.port_sessions' / 'accounts' / 'alice' / 'sessions' / 's1' / 'scratchpad'
scratchpad.mkdir(parents=True)
runtime = PlanRuntime.from_storage_path(scratchpad / 'plan_runtime.json')
runtime.update_plan(
[{'step': 'Inspect runtime planning', 'status': 'pending'}],
explanation='Track the current plan.',
)
runtime_config = AgentRuntimeConfig(cwd=workspace)
model_config = ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct')
prompt_context = build_prompt_context(runtime_config, model_config)
prompt_context = build_prompt_context(
runtime_config,
model_config,
scratchpad_directory=scratchpad,
)
parts = build_system_prompt_parts(
prompt_context=prompt_context,
runtime_config=runtime_config,
tests/test_plan_runtime.py
+5 -1
View File
@@ -183,7 +183,11 @@ class PlanRuntimeTests(unittest.TestCase):
),
)
result = agent.run('Store the current plan')
self.assertTrue((workspace / '.port_sessions' / 'plan_runtime.json').exists())
self.assertIsNotNone(result.scratchpad_directory)
self.assertTrue(
(Path(result.scratchpad_directory) / 'plan_runtime.json').exists()
)
self.assertFalse((workspace / '.port_sessions' / 'plan_runtime.json').exists())
self.assertEqual(result.final_output, 'The plan was stored successfully.')
self.assertEqual(result.tool_calls, 1)
tests/test_task_runtime.py
+5 -1
View File
@@ -233,7 +233,11 @@ class TaskRuntimeTests(unittest.TestCase):
),
)
result = agent.run('Create a task for the current work')
self.assertTrue((workspace / '.port_sessions' / 'task_runtime.json').exists())
self.assertIsNotNone(result.scratchpad_directory)
self.assertTrue(
(Path(result.scratchpad_directory) / 'task_runtime.json').exists()
)
self.assertFalse((workspace / '.port_sessions' / 'task_runtime.json').exists())
self.assertEqual(result.final_output, 'The task was created successfully.')
self.assertEqual(result.tool_calls, 1)