Skip to content

Commit b99643b

Browse files
SkobeltsynSkobeltsyn
authored
Merge pull request #14 from Deep-CodeAI/fix/auth-secret-hygiene
Fix/auth secret hygiene
2 parents 2c23e82 + b9855af commit b99643b

5 files changed

Lines changed: 208 additions & 6 deletions

File tree

‎src/main/kotlin/agents_engine/core/Skill.kt‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,21 @@ class Skill<IN, OUT>(
8282
outputTransformer = block
8383
}
8484

85+
/**
86+
* #856 — opt-in for memory tools. When ANY skill on an agent calls `useMemory()`,
87+
* the agentic loop respects the opt-in: only skills that called this get
88+
* `memory_read` / `memory_write` / `memory_search` in their allowlist. When NO
89+
* skill opts in, the legacy auto-inject (every skill gets memory if memoryBank
90+
* is set) is preserved for backward compatibility.
91+
*/
92+
var useMemory: Boolean = false
93+
private set
94+
95+
fun useMemory() {
96+
checkNotFrozen()
97+
useMemory = true
98+
}
99+
85100
fun execute(input: IN): OUT {
86101
val impl = checkNotNull(implementation) {
87102
"Skill \"$name\" has no implementation. Add implementedBy { } block."

‎src/main/kotlin/agents_engine/mcp/McpAuth.kt‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,19 @@ sealed interface McpAuth {
99
/** No credentials sent. Default. */
1010
object None : McpAuth
1111

12-
/** Sends `Authorization: Bearer <token>` on every request. */
13-
data class Bearer(val token: String) : McpAuth
12+
/**
13+
* Sends `Authorization: Bearer <token>` on every request.
14+
*
15+
* The token is held as a `String` for compatibility with `data class`
16+
* equality. **toString() is overridden to redact the token** so a
17+
* `println(auth)` or `Throwable.message` referencing this instance
18+
* doesn't leak the credential into logs (#857).
19+
*
20+
* Future work: store the token as `CharArray` and add an explicit
21+
* `close()` that wipes it. That requires a richer API surface (the
22+
* transport must know when to wipe) and is left as a follow-up.
23+
*/
24+
data class Bearer(val token: String) : McpAuth {
25+
override fun toString(): String = "Bearer(token=<redacted>)"
26+
}
1427
}

‎src/main/kotlin/agents_engine/model/AgenticLoop.kt‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,12 +28,19 @@ suspend fun <IN> executeAgentic(
2828

2929
val messages = mutableListOf<LlmMessage>()
3030

31-
// Action tools: tools the skill explicitly lists + agent capabilities + auto-injected memory tools
31+
// Action tools: tools the skill explicitly lists + agent capabilities + memory tools
3232
val skillToolDefs = skill.toolNames?.mapNotNull { agent.toolMap[it] } ?: emptyList()
3333
val autoToolDefs = agent.autoToolNames.mapNotNull { agent.toolMap[it] }
34-
val memoryToolDefs = if (agent.memoryBank != null)
35-
agent.toolMap.values.filter { it.name in setOf("memory_read", "memory_write", "memory_search") }
36-
else emptyList()
34+
// #856 — memory-tool authorization is per-skill when ANY skill opts in via
35+
// `useMemory()`. If none opt in, fall back to the legacy default-on behavior
36+
// (every skill gets memory tools when memoryBank is set) so existing
37+
// single-skill agents don't break.
38+
val anySkillOptedIntoMemory = agent.skills.values.any { it.useMemory }
39+
val memoryToolDefs = when {
40+
agent.memoryBank == null -> emptyList()
41+
anySkillOptedIntoMemory && !skill.useMemory -> emptyList()
42+
else -> agent.toolMap.values.filter { it.name in setOf("memory_read", "memory_write", "memory_search") }
43+
}
3744
val actionToolDefs = (skillToolDefs + autoToolDefs + memoryToolDefs).distinctBy { it.name }
3845

3946
// Knowledge tools: exposed lazily — LLM calls them to load context on demand

‎src/test/kotlin/agents_engine/mcp/McpAuthRedactionTest.kt‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package agents_engine.mcp
2+
3+
import kotlin.test.Test
4+
import kotlin.test.assertEquals
5+
import kotlin.test.assertFalse
6+
import kotlin.test.assertTrue
7+
8+
// Tests for #857 — McpAuth.Bearer.toString() must NOT leak the token.
9+
class McpAuthRedactionTest {
10+
11+
@Test
12+
fun `Bearer toString does not contain the token value`() {
13+
val secret = "sk-very-secret-token-do-not-leak"
14+
val auth = McpAuth.Bearer(secret)
15+
val rendered = auth.toString()
16+
assertFalse(rendered.contains(secret), "toString must not include the raw token; got: '$rendered'")
17+
assertTrue(rendered.contains("redacted", ignoreCase = true), "toString must indicate redaction; got: '$rendered'")
18+
}
19+
20+
@Test
21+
fun `Bearer toString does not leak via interpolation`() {
22+
val secret = "sk-another-secret-token"
23+
val auth = McpAuth.Bearer(secret)
24+
val s = "$auth"
25+
assertFalse(s.contains(secret), "string interpolation must not leak token; got: '$s'")
26+
}
27+
28+
@Test
29+
fun `Bearer equals and hashCode still work (data-class semantics preserved)`() {
30+
val a = McpAuth.Bearer("same")
31+
val b = McpAuth.Bearer("same")
32+
val c = McpAuth.Bearer("different")
33+
assertEquals(a, b)
34+
assertEquals(a.hashCode(), b.hashCode())
35+
assertFalse(a == c, "different tokens must not be equal")
36+
}
37+
38+
@Test
39+
fun `Bearer token property still returns the actual value (programmatic access preserved)`() {
40+
val secret = "sk-programmatic"
41+
val auth = McpAuth.Bearer(secret)
42+
assertEquals(secret, auth.token, "the .token property must still expose the value for HTTP-header building")
43+
}
44+
}

‎src/test/kotlin/agents_engine/model/MemoryToolPerSkillOptInTest.kt‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
package agents_engine.model
2+
3+
import agents_engine.core.MemoryBank
4+
import agents_engine.core.agent
5+
import kotlin.test.Test
6+
import kotlin.test.assertEquals
7+
import kotlin.test.assertFalse
8+
import kotlin.test.assertTrue
9+
10+
// Tests for #856 — per-skill memory-tool opt-in.
11+
//
12+
// When ANY skill on an agent calls `useMemory()`, the agentic loop respects
13+
// that opt-in: skills that did NOT opt in must NOT receive memory_* tools in
14+
// their authorized set. When no skill opts in, the legacy default-on behavior
15+
// is preserved for single-skill agents that already work today.
16+
class MemoryToolPerSkillOptInTest {
17+
18+
private fun snoopAuthorizedTools(captureLatest: MutableList<List<String>>): ModelClient {
19+
// Use the system message that the agentic loop builds — it lists every tool
20+
// available to the current skill. We can read it back from the captured
21+
// messages of the first chat() call.
22+
return ModelClient { msgs ->
23+
val sys = msgs.firstOrNull { it.role == "system" }
24+
val tools = sys?.content
25+
?.lineSequence()
26+
?.filter { it.startsWith("- ") }
27+
?.map { it.removePrefix("- ").substringBefore(":") }
28+
?.toList()
29+
?: emptyList()
30+
captureLatest.add(tools)
31+
LlmResponse.Text("done")
32+
}
33+
}
34+
35+
@Test
36+
fun `legacy single-skill agent with memoryBank still auto-gets memory tools (backward compat)`() {
37+
val captured = mutableListOf<List<String>>()
38+
val a = agent<String, String>("legacy") {
39+
model { ollama("test"); client = snoopAuthorizedTools(captured) }
40+
memory(MemoryBank())
41+
skills { skill<String, String>("only", "stub") { tools() } }
42+
}
43+
a("hi")
44+
assertTrue(captured.single().contains("memory_read"), "legacy auto-inject must still work; got: ${captured.single()}")
45+
assertTrue(captured.single().contains("memory_write"))
46+
assertTrue(captured.single().contains("memory_search"))
47+
}
48+
49+
@Test
50+
fun `when one skill opts in, a non-opted-in skill on the same agent loses memory access`() {
51+
// The core security guarantee — auto-inject across ALL skills was the bug.
52+
// With two skills on one agent (one opts in, one doesn't), the non-opted
53+
// skill's authorized tool set must NOT include memory_*.
54+
// The mock client first answers the skill-router LLM call (returning a
55+
// SkillRoute pointing at "read-only"), then snoops the second call's
56+
// system prompt for the "read-only" skill's authorized tools.
57+
val captured = mutableListOf<List<String>>()
58+
var callIdx = 0
59+
val client = ModelClient { msgs ->
60+
callIdx++
61+
if (callIdx == 1) {
62+
// Skill router call — return a JSON SkillRoute.
63+
LlmResponse.Text("""{"skillName":"read-only","confidence":1.0,"rationale":"x"}""")
64+
} else {
65+
// Skill execution call — capture the system prompt's tool listing.
66+
val sys = msgs.firstOrNull { it.role == "system" }
67+
captured.add(
68+
sys?.content?.lineSequence()
69+
?.filter { it.startsWith("- ") }
70+
?.map { it.removePrefix("- ").substringBefore(":") }
71+
?.toList()
72+
?: emptyList(),
73+
)
74+
LlmResponse.Text("done")
75+
}
76+
}
77+
val a = agent<String, String>("two-skill") {
78+
model { ollama("test"); this.client = client }
79+
memory(MemoryBank())
80+
skills {
81+
skill<String, String>("read-only", "answers questions") { tools() }
82+
skill<String, String>("memo-writer", "writes notes") { tools(); useMemory() }
83+
}
84+
}
85+
a("input")
86+
val tools = captured.single()
87+
assertFalse(tools.contains("memory_read"), "non-opted-in skill must NOT receive memory_read; got: $tools")
88+
assertFalse(tools.contains("memory_write"), "non-opted-in skill must NOT receive memory_write; got: $tools")
89+
assertFalse(tools.contains("memory_search"), "non-opted-in skill must NOT receive memory_search; got: $tools")
90+
}
91+
92+
@Test
93+
fun `opted-in skill DOES receive memory tools`() {
94+
val captured = mutableListOf<List<String>>()
95+
val a = agent<String, String>("writer") {
96+
model { ollama("test"); client = snoopAuthorizedTools(captured) }
97+
memory(MemoryBank())
98+
skills {
99+
skill<String, String>("memo-writer", "uses memory") { tools(); useMemory() }
100+
}
101+
}
102+
a("hi")
103+
val tools = captured.single()
104+
assertTrue(tools.contains("memory_read"), "opted-in skill must receive memory_read; got: $tools")
105+
assertTrue(tools.contains("memory_write"))
106+
assertTrue(tools.contains("memory_search"))
107+
}
108+
109+
@Test
110+
fun `agent without memoryBank never injects memory tools regardless of useMemory`() {
111+
val captured = mutableListOf<List<String>>()
112+
val a = agent<String, String>("no-memory") {
113+
model { ollama("test"); client = snoopAuthorizedTools(captured) }
114+
// no memory(...) call
115+
skills {
116+
skill<String, String>("s", "stub") { tools(); useMemory() }
117+
}
118+
}
119+
a("hi")
120+
val tools = captured.single()
121+
assertEquals(emptyList(), tools.filter { it.startsWith("memory_") }, "no memoryBank means no memory tools regardless of opt-in")
122+
}
123+
}

0 commit comments

Comments
 (0)