kaiwen-agent 0.1.0a1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (52) hide show
  1. kaiwen_agent-0.1.0a1/.gitignore +17 -0
  2. kaiwen_agent-0.1.0a1/PKG-INFO +124 -0
  3. kaiwen_agent-0.1.0a1/README.md +91 -0
  4. kaiwen_agent-0.1.0a1/pyproject.toml +66 -0
  5. kaiwen_agent-0.1.0a1/src/kaiwen_agent/__init__.py +18 -0
  6. kaiwen_agent-0.1.0a1/src/kaiwen_agent/agent.py +136 -0
  7. kaiwen_agent-0.1.0a1/src/kaiwen_agent/audit.py +84 -0
  8. kaiwen_agent-0.1.0a1/src/kaiwen_agent/context.py +12 -0
  9. kaiwen_agent-0.1.0a1/src/kaiwen_agent/control/__init__.py +17 -0
  10. kaiwen_agent-0.1.0a1/src/kaiwen_agent/control/checkpoints.py +39 -0
  11. kaiwen_agent-0.1.0a1/src/kaiwen_agent/control/mailbox.py +61 -0
  12. kaiwen_agent-0.1.0a1/src/kaiwen_agent/control/runtime.py +41 -0
  13. kaiwen_agent-0.1.0a1/src/kaiwen_agent/events.py +94 -0
  14. kaiwen_agent-0.1.0a1/src/kaiwen_agent/exceptions.py +34 -0
  15. kaiwen_agent-0.1.0a1/src/kaiwen_agent/models/__init__.py +5 -0
  16. kaiwen_agent-0.1.0a1/src/kaiwen_agent/models/base.py +18 -0
  17. kaiwen_agent-0.1.0a1/src/kaiwen_agent/models/fake.py +27 -0
  18. kaiwen_agent-0.1.0a1/src/kaiwen_agent/models/openai.py +178 -0
  19. kaiwen_agent-0.1.0a1/src/kaiwen_agent/persistence/__init__.py +3 -0
  20. kaiwen_agent-0.1.0a1/src/kaiwen_agent/persistence/sqlite.py +371 -0
  21. kaiwen_agent-0.1.0a1/src/kaiwen_agent/py.typed +1 -0
  22. kaiwen_agent-0.1.0a1/src/kaiwen_agent/resources/__init__.py +3 -0
  23. kaiwen_agent-0.1.0a1/src/kaiwen_agent/resources/manager.py +39 -0
  24. kaiwen_agent-0.1.0a1/src/kaiwen_agent/sessions.py +39 -0
  25. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tasks/__init__.py +19 -0
  26. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tasks/graph.py +108 -0
  27. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tasks/models.py +75 -0
  28. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tasks/scheduler.py +259 -0
  29. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tasks/state.py +59 -0
  30. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tasks/store.py +33 -0
  31. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tools/__init__.py +6 -0
  32. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tools/approval.py +29 -0
  33. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tools/decorator.py +39 -0
  34. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tools/definition.py +42 -0
  35. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tools/execution.py +220 -0
  36. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tools/idempotency.py +23 -0
  37. kaiwen_agent-0.1.0a1/src/kaiwen_agent/tools/registry.py +26 -0
  38. kaiwen_agent-0.1.0a1/src/kaiwen_agent/transports/__init__.py +3 -0
  39. kaiwen_agent-0.1.0a1/src/kaiwen_agent/transports/sse.py +52 -0
  40. kaiwen_agent-0.1.0a1/src/kaiwen_agent/types.py +76 -0
  41. kaiwen_agent-0.1.0a1/src/kaiwen_agent/workers/__init__.py +16 -0
  42. kaiwen_agent-0.1.0a1/src/kaiwen_agent/workers/base.py +60 -0
  43. kaiwen_agent-0.1.0a1/src/kaiwen_agent/workers/registry.py +21 -0
  44. kaiwen_agent-0.1.0a1/tests/__init__.py +1 -0
  45. kaiwen_agent-0.1.0a1/tests/test_agent.py +120 -0
  46. kaiwen_agent-0.1.0a1/tests/test_openai_provider.py +162 -0
  47. kaiwen_agent-0.1.0a1/tests/test_persistence.py +136 -0
  48. kaiwen_agent-0.1.0a1/tests/test_protocol_contract.py +25 -0
  49. kaiwen_agent-0.1.0a1/tests/test_sse.py +42 -0
  50. kaiwen_agent-0.1.0a1/tests/test_tasks.py +192 -0
  51. kaiwen_agent-0.1.0a1/tests/test_tool_pipeline.py +152 -0
  52. kaiwen_agent-0.1.0a1/tests/test_tools.py +18 -0
@@ -0,0 +1,17 @@
1
+ .venv/
2
+ __pycache__/
3
+ .pytest_cache/
4
+ .ruff_cache/
5
+ .mypy_cache/
6
+ *.py[cod]
7
+ dist/
8
+ build/
9
+ *.egg-info/
10
+ node_modules/
11
+ coverage/
12
+ .env
13
+ .env.*
14
+ !.env.example
15
+ *.sqlite3
16
+ *.sqlite3-shm
17
+ *.sqlite3-wal
@@ -0,0 +1,124 @@
1
+ Metadata-Version: 2.4
2
+ Name: kaiwen-agent
3
+ Version: 0.1.0a1
4
+ Summary: Provider-neutral asynchronous agent runtime
5
+ Project-URL: Homepage, https://github.com/ImtheKaiwen/kaiwen-agent-platform
6
+ Project-URL: Documentation, https://github.com/ImtheKaiwen/kaiwen-agent-platform/tree/main/docs
7
+ Project-URL: Repository, https://github.com/ImtheKaiwen/kaiwen-agent-platform.git
8
+ Project-URL: Issues, https://github.com/ImtheKaiwen/kaiwen-agent-platform/issues
9
+ Author: Kaiwen
10
+ License: MIT
11
+ Keywords: agent,ai,asyncio,openai,tools
12
+ Classifier: Development Status :: 3 - Alpha
13
+ Classifier: License :: OSI Approved :: MIT License
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.11
16
+ Classifier: Programming Language :: Python :: 3.12
17
+ Classifier: Programming Language :: Python :: 3.13
18
+ Classifier: Typing :: Typed
19
+ Requires-Python: >=3.11
20
+ Requires-Dist: pydantic<3,>=2.10
21
+ Provides-Extra: dev
22
+ Requires-Dist: mypy<2,>=1.14; extra == 'dev'
23
+ Requires-Dist: openai<3,>=2; extra == 'dev'
24
+ Requires-Dist: pytest<9,>=8.3; extra == 'dev'
25
+ Requires-Dist: python-dotenv<2,>=1.1; extra == 'dev'
26
+ Requires-Dist: ruff<1,>=0.9; extra == 'dev'
27
+ Provides-Extra: openai
28
+ Requires-Dist: openai<3,>=2; extra == 'openai'
29
+ Provides-Extra: release
30
+ Requires-Dist: build<2,>=1.2; extra == 'release'
31
+ Requires-Dist: twine<7,>=6; extra == 'release'
32
+ Description-Content-Type: text/markdown
33
+
34
+ # kaiwen-agent
35
+
36
+ Provider-neutral, asynchronous Python runtime for building tool-using agents.
37
+ Product applications own their prompts, domain tools, authorization context,
38
+ persistence choice, and HTTP surface.
39
+
40
+ ```bash
41
+ pip install "kaiwen-agent[openai]"
42
+ ```
43
+
44
+ ## Minimal agent
45
+
46
+ ```python
47
+ import asyncio
48
+
49
+ from pydantic import BaseModel
50
+
51
+ from kaiwen_agent import Agent, AgentContext, tool
52
+ from kaiwen_agent.models.openai import OpenAIResponsesProvider
53
+ from kaiwen_agent.types import AgentInput
54
+
55
+
56
+ class ProjectQuery(BaseModel):
57
+ category: str | None = None
58
+
59
+
60
+ @tool(
61
+ name="projects.read",
62
+ description="List projects visible to the current user",
63
+ input_model=ProjectQuery,
64
+ permissions={"projects.read"},
65
+ )
66
+ async def read_projects(
67
+ context: AgentContext,
68
+ query: ProjectQuery,
69
+ ) -> list[dict[str, str]]:
70
+ repository = context.services["projects"]
71
+ return await repository.list(category=query.category)
72
+
73
+
74
+ async def main() -> None:
75
+ agent = Agent(
76
+ name="portfolio",
77
+ model=OpenAIResponsesProvider(
78
+ model="your-model-id",
79
+ instructions="Help visitors discover relevant products.",
80
+ store=False,
81
+ ),
82
+ tools=[read_projects],
83
+ permissions={"projects.read"},
84
+ )
85
+ result = await agent.run(
86
+ AgentInput(text="Show me the games"),
87
+ context=AgentContext(
88
+ session_id="session-id",
89
+ permissions=frozenset({"projects.read"}),
90
+ services={"projects": your_repository},
91
+ ),
92
+ )
93
+ print(result.text)
94
+
95
+
96
+ asyncio.run(main())
97
+ ```
98
+
99
+ `OPENAI_API_KEY` is read by the OpenAI SDK from the server environment. The
100
+ runtime never stores it and it must never be exposed to browser code.
101
+
102
+ ## Runtime capabilities
103
+
104
+ - Provider-neutral async agent loop and deterministic fake provider.
105
+ - Typed Pydantic tool inputs and deny-by-default permission intersection.
106
+ - Approval, idempotency, timeout, and retry handling for side effects.
107
+ - Session, run, event, and audit ports with a SQLite reference adapter.
108
+ - SSE event fan-out and trace propagation.
109
+ - Durable task graphs, workers, checkpoints, mailbox control, resource locks,
110
+ cancellation, and bounded concurrency.
111
+ - Optional OpenAI Responses API adapter.
112
+
113
+ ## Side effects
114
+
115
+ Tools marked `reversible` or `destructive` receive a runtime idempotency key.
116
+ Set `requires_approval=True` and supply an application-owned `ApprovalGate` for
117
+ publishing, deleting, deploying, calling, or sending messages. Applications
118
+ must still perform authorization inside their own service boundary.
119
+
120
+ ## Coding-agent handoff
121
+
122
+ When asking another coding agent to integrate this package, provide the root
123
+ repository files `AGENTS.md` and `docs/consumer-agent-brief.md`. The complete,
124
+ provider-free example is in `examples/basic-agent`.
@@ -0,0 +1,91 @@
1
+ # kaiwen-agent
2
+
3
+ Provider-neutral, asynchronous Python runtime for building tool-using agents.
4
+ Product applications own their prompts, domain tools, authorization context,
5
+ persistence choice, and HTTP surface.
6
+
7
+ ```bash
8
+ pip install "kaiwen-agent[openai]"
9
+ ```
10
+
11
+ ## Minimal agent
12
+
13
+ ```python
14
+ import asyncio
15
+
16
+ from pydantic import BaseModel
17
+
18
+ from kaiwen_agent import Agent, AgentContext, tool
19
+ from kaiwen_agent.models.openai import OpenAIResponsesProvider
20
+ from kaiwen_agent.types import AgentInput
21
+
22
+
23
+ class ProjectQuery(BaseModel):
24
+ category: str | None = None
25
+
26
+
27
+ @tool(
28
+ name="projects.read",
29
+ description="List projects visible to the current user",
30
+ input_model=ProjectQuery,
31
+ permissions={"projects.read"},
32
+ )
33
+ async def read_projects(
34
+ context: AgentContext,
35
+ query: ProjectQuery,
36
+ ) -> list[dict[str, str]]:
37
+ repository = context.services["projects"]
38
+ return await repository.list(category=query.category)
39
+
40
+
41
+ async def main() -> None:
42
+ agent = Agent(
43
+ name="portfolio",
44
+ model=OpenAIResponsesProvider(
45
+ model="your-model-id",
46
+ instructions="Help visitors discover relevant products.",
47
+ store=False,
48
+ ),
49
+ tools=[read_projects],
50
+ permissions={"projects.read"},
51
+ )
52
+ result = await agent.run(
53
+ AgentInput(text="Show me the games"),
54
+ context=AgentContext(
55
+ session_id="session-id",
56
+ permissions=frozenset({"projects.read"}),
57
+ services={"projects": your_repository},
58
+ ),
59
+ )
60
+ print(result.text)
61
+
62
+
63
+ asyncio.run(main())
64
+ ```
65
+
66
+ `OPENAI_API_KEY` is read by the OpenAI SDK from the server environment. The
67
+ runtime never stores it and it must never be exposed to browser code.
68
+
69
+ ## Runtime capabilities
70
+
71
+ - Provider-neutral async agent loop and deterministic fake provider.
72
+ - Typed Pydantic tool inputs and deny-by-default permission intersection.
73
+ - Approval, idempotency, timeout, and retry handling for side effects.
74
+ - Session, run, event, and audit ports with a SQLite reference adapter.
75
+ - SSE event fan-out and trace propagation.
76
+ - Durable task graphs, workers, checkpoints, mailbox control, resource locks,
77
+ cancellation, and bounded concurrency.
78
+ - Optional OpenAI Responses API adapter.
79
+
80
+ ## Side effects
81
+
82
+ Tools marked `reversible` or `destructive` receive a runtime idempotency key.
83
+ Set `requires_approval=True` and supply an application-owned `ApprovalGate` for
84
+ publishing, deleting, deploying, calling, or sending messages. Applications
85
+ must still perform authorization inside their own service boundary.
86
+
87
+ ## Coding-agent handoff
88
+
89
+ When asking another coding agent to integrate this package, provide the root
90
+ repository files `AGENTS.md` and `docs/consumer-agent-brief.md`. The complete,
91
+ provider-free example is in `examples/basic-agent`.
@@ -0,0 +1,66 @@
1
+ [build-system]
2
+ requires = ["hatchling>=1.27,<1.32"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "kaiwen-agent"
7
+ version = "0.1.0a1"
8
+ description = "Provider-neutral asynchronous agent runtime"
9
+ readme = "README.md"
10
+ requires-python = ">=3.11"
11
+ license = { text = "MIT" }
12
+ authors = [{ name = "Kaiwen" }]
13
+ keywords = ["agent", "ai", "asyncio", "openai", "tools"]
14
+ classifiers = [
15
+ "Development Status :: 3 - Alpha",
16
+ "License :: OSI Approved :: MIT License",
17
+ "Programming Language :: Python :: 3",
18
+ "Programming Language :: Python :: 3.11",
19
+ "Programming Language :: Python :: 3.12",
20
+ "Programming Language :: Python :: 3.13",
21
+ "Typing :: Typed",
22
+ ]
23
+ dependencies = [
24
+ "pydantic>=2.10,<3",
25
+ ]
26
+
27
+ [project.urls]
28
+ Homepage = "https://github.com/ImtheKaiwen/kaiwen-agent-platform"
29
+ Documentation = "https://github.com/ImtheKaiwen/kaiwen-agent-platform/tree/main/docs"
30
+ Repository = "https://github.com/ImtheKaiwen/kaiwen-agent-platform.git"
31
+ Issues = "https://github.com/ImtheKaiwen/kaiwen-agent-platform/issues"
32
+
33
+ [project.optional-dependencies]
34
+ openai = [
35
+ "openai>=2,<3",
36
+ ]
37
+ dev = [
38
+ "mypy>=1.14,<2",
39
+ "python-dotenv>=1.1,<2",
40
+ "pytest>=8.3,<9",
41
+ "ruff>=0.9,<1",
42
+ "openai>=2,<3",
43
+ ]
44
+ release = [
45
+ "build>=1.2,<2",
46
+ "twine>=6,<7",
47
+ ]
48
+
49
+ [tool.hatch.build.targets.wheel]
50
+ packages = ["src/kaiwen_agent"]
51
+
52
+ [tool.pytest.ini_options]
53
+ addopts = "-q -p no:cacheprovider"
54
+ testpaths = ["tests"]
55
+
56
+ [tool.ruff]
57
+ line-length = 100
58
+ target-version = "py311"
59
+
60
+ [tool.ruff.lint]
61
+ select = ["E", "F", "I", "UP", "B", "SIM"]
62
+
63
+ [tool.mypy]
64
+ python_version = "3.11"
65
+ strict = true
66
+ packages = ["kaiwen_agent"]
@@ -0,0 +1,18 @@
1
+ from kaiwen_agent.agent import Agent
2
+ from kaiwen_agent.context import AgentContext
3
+ from kaiwen_agent.models.base import ModelProvider
4
+ from kaiwen_agent.tools.approval import ApprovalGate, ApprovalRequest
5
+ from kaiwen_agent.tools.decorator import tool
6
+ from kaiwen_agent.tools.definition import RetryPolicy
7
+ from kaiwen_agent.tools.execution import ToolExecutor
8
+
9
+ __all__ = [
10
+ "Agent",
11
+ "AgentContext",
12
+ "ApprovalGate",
13
+ "ApprovalRequest",
14
+ "ModelProvider",
15
+ "RetryPolicy",
16
+ "ToolExecutor",
17
+ "tool",
18
+ ]
@@ -0,0 +1,136 @@
1
+ from __future__ import annotations
2
+
3
+ from datetime import UTC, datetime
4
+ from uuid import uuid4
5
+
6
+ from kaiwen_agent.context import AgentContext
7
+ from kaiwen_agent.events import AgentEvent, EventSink, EventType, InMemoryEventSink
8
+ from kaiwen_agent.models.base import ModelProvider
9
+ from kaiwen_agent.sessions import InMemoryRunStore, RunStore
10
+ from kaiwen_agent.tools.definition import ToolDefinition
11
+ from kaiwen_agent.tools.execution import ToolExecutor
12
+ from kaiwen_agent.tools.registry import ToolRegistry
13
+ from kaiwen_agent.types import AgentInput, AgentResult, AgentRun, ModelUsage, ToolResult
14
+
15
+
16
+ class Agent:
17
+ def __init__(
18
+ self,
19
+ *,
20
+ name: str,
21
+ model: ModelProvider,
22
+ tools: list[ToolDefinition] | None = None,
23
+ permissions: set[str] | frozenset[str] | None = None,
24
+ event_sink: EventSink | None = None,
25
+ run_store: RunStore | None = None,
26
+ tool_executor: ToolExecutor | None = None,
27
+ max_steps: int = 8,
28
+ ) -> None:
29
+ if max_steps < 1:
30
+ raise ValueError("max_steps must be at least 1")
31
+ self.name = name
32
+ self.model = model
33
+ self.registry = ToolRegistry(tools or [])
34
+ self.permissions = frozenset(permissions or set())
35
+ self.event_sink = event_sink or InMemoryEventSink()
36
+ self.run_store = run_store or InMemoryRunStore()
37
+ self.tool_executor = tool_executor or ToolExecutor()
38
+ self.max_steps = max_steps
39
+
40
+ async def run(self, agent_input: AgentInput, *, context: AgentContext) -> AgentResult:
41
+ trace_id = context.trace_id or f"trace_{uuid4().hex}"
42
+ effective_context = AgentContext(
43
+ session_id=context.session_id,
44
+ trace_id=trace_id,
45
+ permissions=context.permissions & self.permissions,
46
+ services=context.services,
47
+ metadata=context.metadata,
48
+ )
49
+ run = AgentRun(session_id=context.session_id, status="running")
50
+ await self.run_store.save(run)
51
+ await self._emit("run.started", run, trace_id, {"agent": self.name})
52
+ tool_results: list[ToolResult] = []
53
+ pending_tool_results: list[ToolResult] = []
54
+ previous_response_id: str | None = None
55
+ total_usage = ModelUsage()
56
+
57
+ try:
58
+ for step in range(self.max_steps):
59
+ response = await self.model.respond(
60
+ agent_input,
61
+ context=effective_context,
62
+ tools=self.registry.definitions(),
63
+ tool_results=pending_tool_results,
64
+ previous_response_id=previous_response_id,
65
+ )
66
+ pending_tool_results = []
67
+ previous_response_id = response.provider_response_id or previous_response_id
68
+ if response.usage is not None:
69
+ total_usage.input_tokens += response.usage.input_tokens
70
+ total_usage.output_tokens += response.usage.output_tokens
71
+ total_usage.total_tokens += response.usage.total_tokens
72
+ await self._emit(
73
+ "run.status",
74
+ run,
75
+ trace_id,
76
+ {
77
+ "kind": "model_usage",
78
+ "model": response.model or "unknown",
79
+ **response.usage.model_dump(),
80
+ },
81
+ )
82
+ if not response.tool_calls:
83
+ result = AgentResult(
84
+ run_id=run.id,
85
+ session_id=run.session_id,
86
+ text=response.text or "",
87
+ status="completed",
88
+ tool_results=tool_results,
89
+ structured_output=response.structured_output,
90
+ usage=total_usage,
91
+ )
92
+ run.status = "completed"
93
+ run.completed_at = datetime.now(UTC)
94
+ await self.run_store.save(run)
95
+ await self._emit(
96
+ "run.completed", run, trace_id, {"step": step, "text": result.text}
97
+ )
98
+ return result
99
+
100
+ for call in response.tool_calls:
101
+ definition = self.registry.get(call.name)
102
+ tool_result = await self.tool_executor.execute(
103
+ definition,
104
+ call,
105
+ effective_context,
106
+ execution_scope=run.id,
107
+ emit=lambda event_type, payload: self._emit(
108
+ event_type, run, trace_id, payload
109
+ ),
110
+ )
111
+ tool_results.append(tool_result)
112
+ pending_tool_results.append(tool_result)
113
+
114
+ raise RuntimeError(f"Agent exceeded max_steps={self.max_steps}")
115
+ except Exception as error:
116
+ run.status = "failed"
117
+ run.completed_at = datetime.now(UTC)
118
+ await self.run_store.save(run)
119
+ await self._emit("run.failed", run, trace_id, {"error_type": type(error).__name__})
120
+ raise
121
+
122
+ async def _emit(
123
+ self,
124
+ event_type: EventType,
125
+ run: AgentRun,
126
+ trace_id: str,
127
+ payload: dict[str, object],
128
+ ) -> None:
129
+ await self.event_sink.append(
130
+ AgentEvent(
131
+ type=event_type,
132
+ trace_id=trace_id,
133
+ session_id=run.session_id,
134
+ payload=payload,
135
+ )
136
+ )
@@ -0,0 +1,84 @@
1
+ from __future__ import annotations
2
+
3
+ import re
4
+ from datetime import datetime
5
+ from typing import Any, Literal, Protocol
6
+ from uuid import uuid4
7
+
8
+ from pydantic import BaseModel, ConfigDict, Field
9
+
10
+ from kaiwen_agent.types import utc_now
11
+
12
+
13
+ class AuditRecord(BaseModel):
14
+ model_config = ConfigDict(extra="forbid")
15
+
16
+ id: str = Field(default_factory=lambda: f"audit_{uuid4().hex}")
17
+ trace_id: str
18
+ session_id: str
19
+ action: str
20
+ outcome: Literal["succeeded", "failed", "denied"]
21
+ details: dict[str, Any] = Field(default_factory=dict)
22
+ timestamp: datetime = Field(default_factory=utc_now)
23
+
24
+
25
+ class AuditStore(Protocol):
26
+ async def append_audit(self, record: AuditRecord) -> None: ...
27
+
28
+ async def list_audit(self, *, trace_id: str | None = None) -> list[AuditRecord]: ...
29
+
30
+
31
+ class InMemoryAuditStore:
32
+ def __init__(self) -> None:
33
+ self.records: list[AuditRecord] = []
34
+
35
+ async def append_audit(self, record: AuditRecord) -> None:
36
+ self.records.append(record.model_copy(deep=True))
37
+
38
+ async def list_audit(self, *, trace_id: str | None = None) -> list[AuditRecord]:
39
+ records = self.records
40
+ if trace_id is not None:
41
+ records = [record for record in records if record.trace_id == trace_id]
42
+ return [record.model_copy(deep=True) for record in records]
43
+
44
+
45
+ class SecretRedactor:
46
+ _secret_keys = re.compile(
47
+ r"(^|_)(api_?key|authorization|cookie|password|secret|token)($|_)", re.IGNORECASE
48
+ )
49
+ _secret_values = (
50
+ re.compile(r"\bsk-[A-Za-z0-9_-]{12,}\b"),
51
+ re.compile(r"\bBearer\s+[A-Za-z0-9._~+/-]+=*", re.IGNORECASE),
52
+ )
53
+
54
+ def redact(self, value: Any, *, key: str | None = None) -> Any:
55
+ if key is not None and self._secret_keys.search(key):
56
+ return "[REDACTED]"
57
+ if isinstance(value, dict):
58
+ return {
59
+ item_key: self.redact(item_value, key=str(item_key))
60
+ for item_key, item_value in value.items()
61
+ }
62
+ if isinstance(value, list):
63
+ return [self.redact(item) for item in value]
64
+ if isinstance(value, tuple):
65
+ return tuple(self.redact(item) for item in value)
66
+ if isinstance(value, str):
67
+ redacted = value
68
+ for pattern in self._secret_values:
69
+ redacted = pattern.sub("[REDACTED]", redacted)
70
+ return redacted
71
+ return value
72
+
73
+
74
+ class RedactingAuditStore:
75
+ def __init__(self, store: AuditStore, redactor: SecretRedactor | None = None) -> None:
76
+ self.store = store
77
+ self.redactor = redactor or SecretRedactor()
78
+
79
+ async def append_audit(self, record: AuditRecord) -> None:
80
+ clean = record.model_copy(update={"details": self.redactor.redact(record.details)})
81
+ await self.store.append_audit(clean)
82
+
83
+ async def list_audit(self, *, trace_id: str | None = None) -> list[AuditRecord]:
84
+ return await self.store.list_audit(trace_id=trace_id)
@@ -0,0 +1,12 @@
1
+ from collections.abc import Mapping
2
+ from dataclasses import dataclass, field
3
+ from typing import Any
4
+
5
+
6
+ @dataclass(slots=True)
7
+ class AgentContext:
8
+ session_id: str
9
+ trace_id: str | None = None
10
+ permissions: frozenset[str] = field(default_factory=frozenset)
11
+ services: Mapping[str, Any] = field(default_factory=dict)
12
+ metadata: Mapping[str, Any] = field(default_factory=dict)
@@ -0,0 +1,17 @@
1
+ from kaiwen_agent.control.checkpoints import (
2
+ CheckpointStore,
3
+ InMemoryCheckpointStore,
4
+ TaskCheckpoint,
5
+ )
6
+ from kaiwen_agent.control.mailbox import ControlMessage, InMemoryMailbox, TaskCancelledError
7
+ from kaiwen_agent.control.runtime import RunControl
8
+
9
+ __all__ = [
10
+ "CheckpointStore",
11
+ "ControlMessage",
12
+ "InMemoryCheckpointStore",
13
+ "InMemoryMailbox",
14
+ "RunControl",
15
+ "TaskCancelledError",
16
+ "TaskCheckpoint",
17
+ ]
@@ -0,0 +1,39 @@
1
+ from __future__ import annotations
2
+
3
+ from datetime import datetime
4
+ from typing import Any, Protocol
5
+ from uuid import uuid4
6
+
7
+ from pydantic import BaseModel, ConfigDict, Field
8
+
9
+ from kaiwen_agent.types import utc_now
10
+
11
+
12
+ class TaskCheckpoint(BaseModel):
13
+ model_config = ConfigDict(extra="forbid")
14
+
15
+ id: str = Field(default_factory=lambda: f"checkpoint_{uuid4().hex}")
16
+ task_id: str
17
+ sequence: int = Field(ge=1)
18
+ state: dict[str, Any] = Field(default_factory=dict)
19
+ created_at: datetime = Field(default_factory=utc_now)
20
+
21
+
22
+ class CheckpointStore(Protocol):
23
+ async def save_checkpoint(self, checkpoint: TaskCheckpoint) -> None: ...
24
+
25
+ async def latest_checkpoint(self, task_id: str) -> TaskCheckpoint | None: ...
26
+
27
+
28
+ class InMemoryCheckpointStore:
29
+ def __init__(self) -> None:
30
+ self._checkpoints: dict[str, list[TaskCheckpoint]] = {}
31
+
32
+ async def save_checkpoint(self, checkpoint: TaskCheckpoint) -> None:
33
+ self._checkpoints.setdefault(checkpoint.task_id, []).append(
34
+ checkpoint.model_copy(deep=True)
35
+ )
36
+
37
+ async def latest_checkpoint(self, task_id: str) -> TaskCheckpoint | None:
38
+ checkpoints = self._checkpoints.get(task_id, [])
39
+ return checkpoints[-1].model_copy(deep=True) if checkpoints else None
@@ -0,0 +1,61 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ from collections import defaultdict, deque
5
+ from datetime import datetime
6
+ from typing import Any, Literal
7
+ from uuid import uuid4
8
+
9
+ from pydantic import BaseModel, ConfigDict, Field
10
+
11
+ from kaiwen_agent.types import utc_now
12
+
13
+ ControlType = Literal["instruction", "cancel", "pause", "resume", "approval"]
14
+
15
+
16
+ class ControlMessage(BaseModel):
17
+ model_config = ConfigDict(extra="forbid")
18
+
19
+ id: str = Field(default_factory=lambda: f"control_{uuid4().hex}")
20
+ type: ControlType
21
+ content: str = ""
22
+ created_by: str
23
+ created_at: datetime = Field(default_factory=utc_now)
24
+ metadata: dict[str, Any] = Field(default_factory=dict)
25
+
26
+
27
+ class InMemoryMailbox:
28
+ def __init__(self) -> None:
29
+ self._messages: dict[str, deque[ControlMessage]] = defaultdict(deque)
30
+ self._conditions: dict[str, asyncio.Condition] = defaultdict(asyncio.Condition)
31
+ self._paused: set[str] = set()
32
+ self._cancelled: set[str] = set()
33
+
34
+ async def send(self, task_id: str, message: ControlMessage) -> None:
35
+ condition = self._conditions[task_id]
36
+ async with condition:
37
+ self._messages[task_id].append(message.model_copy(deep=True))
38
+ if message.type == "cancel":
39
+ self._cancelled.add(task_id)
40
+ elif message.type == "pause":
41
+ self._paused.add(task_id)
42
+ elif message.type == "resume":
43
+ self._paused.discard(task_id)
44
+ condition.notify_all()
45
+
46
+ async def checkpoint(self, task_id: str) -> list[ControlMessage]:
47
+ condition = self._conditions[task_id]
48
+ async with condition:
49
+ if task_id in self._cancelled:
50
+ raise TaskCancelledError(f"Task cancelled: {task_id}")
51
+ while task_id in self._paused:
52
+ await condition.wait()
53
+ if task_id in self._cancelled:
54
+ raise TaskCancelledError(f"Task cancelled: {task_id}")
55
+ messages = list(self._messages[task_id])
56
+ self._messages[task_id].clear()
57
+ return [message.model_copy(deep=True) for message in messages]
58
+
59
+
60
+ class TaskCancelledError(Exception):
61
+ pass